mockforge_intelligence/intelligent_behavior/
llm_client.rs

1//! LLM client wrapper for intelligent behavior
2//!
3//! This module provides a simplified interface to the RAG engine for
4//! intelligent mock behavior generation.
5
6use std::sync::Arc;
7use tokio::sync::RwLock;
8
9use super::config::BehaviorModelConfig;
10use super::types::LlmGenerationRequest;
11use mockforge_foundation::Result;
12
13/// LLM client for generating intelligent responses
14pub struct LlmClient {
15    /// RAG engine (lazily initialized)
16    rag_engine: Arc<RwLock<Option<Box<dyn LlmProvider>>>>,
17    /// Configuration
18    config: BehaviorModelConfig,
19}
20
21impl LlmClient {
22    /// Create a new LLM client
23    pub fn new(config: BehaviorModelConfig) -> Self {
24        Self {
25            rag_engine: Arc::new(RwLock::new(None)),
26            config,
27        }
28    }
29
30    /// Initialize the RAG engine (lazy initialization)
31    async fn ensure_initialized(&self) -> Result<()> {
32        let mut engine = self.rag_engine.write().await;
33
34        if engine.is_none() {
35            // Create provider based on configuration
36            let provider = self.create_provider()?;
37            *engine = Some(provider);
38        }
39
40        Ok(())
41    }
42
43    /// Create LLM provider based on configuration
44    fn create_provider(&self) -> Result<Box<dyn LlmProvider>> {
45        match self.config.llm_provider.to_lowercase().as_str() {
46            "openai" => Ok(Box::new(OpenAIProvider::new(&self.config)?)),
47            "anthropic" => Ok(Box::new(AnthropicProvider::new(&self.config)?)),
48            "ollama" => Ok(Box::new(OllamaProvider::new(&self.config)?)),
49            "openai-compatible" => Ok(Box::new(OpenAICompatibleProvider::new(&self.config)?)),
50            _ => Err(mockforge_foundation::Error::internal(format!(
51                "Unsupported LLM provider: {}",
52                self.config.llm_provider
53            ))),
54        }
55    }
56
57    /// Resolve the effective sampling seed (#852): per-request override,
58    /// then the behavior-model config, then the MOCKFORGE_AI_SEED env var.
59    fn resolve_seed(&self, request: &LlmGenerationRequest) -> Option<i64> {
60        request
61            .seed
62            .or(self.config.seed)
63            .or_else(|| std::env::var("MOCKFORGE_AI_SEED").ok().and_then(|s| s.trim().parse().ok()))
64    }
65
66    /// Generate a response and require one valid JSON value without fallback wrappers.
67    pub async fn generate_json(&self, request: &LlmGenerationRequest) -> Result<serde_json::Value> {
68        let text = self.generate_text(request).await?;
69        serde_json::from_str(&text).map_err(|_| {
70            mockforge_foundation::Error::internal(
71                "AI provider returned invalid JSON; no fixture was baked",
72            )
73        })
74    }
75
76    async fn generate_text(&self, request: &LlmGenerationRequest) -> Result<String> {
77        self.ensure_initialized().await?;
78
79        let engine = self.rag_engine.read().await;
80        let provider = engine
81            .as_ref()
82            .ok_or_else(|| mockforge_foundation::Error::internal("LLM provider not initialized"))?;
83
84        // Build messages
85        let messages = vec![
86            ChatMessage {
87                role: "system".to_string(),
88                content: request.system_prompt.clone(),
89            },
90            ChatMessage {
91                role: "user".to_string(),
92                content: request.user_prompt.clone(),
93            },
94        ];
95
96        // Generate response
97        provider
98            .generate_chat(
99                messages,
100                request.temperature,
101                request.max_tokens,
102                self.resolve_seed(request),
103            )
104            .await
105    }
106
107    /// Generate a response from a prompt
108    pub async fn generate(&self, request: &LlmGenerationRequest) -> Result<serde_json::Value> {
109        let response_text = self.generate_text(request).await?;
110
111        // Try to parse as JSON
112        match serde_json::from_str::<serde_json::Value>(&response_text) {
113            Ok(json) => Ok(json),
114            Err(_) => {
115                // Try to extract JSON from response
116                if let Some(start) = response_text.find('{') {
117                    if let Some(end) = response_text.rfind('}') {
118                        let json_str = &response_text[start..=end];
119                        if let Ok(json) = serde_json::from_str::<serde_json::Value>(json_str) {
120                            return Ok(json);
121                        }
122                    }
123                }
124
125                // Fallback: wrap in object
126                Ok(serde_json::json!({
127                    "response": response_text,
128                    "note": "Response was not valid JSON, wrapped in object"
129                }))
130            }
131        }
132    }
133
134    /// Generate a response and return usage information
135    pub async fn generate_with_usage(
136        &self,
137        request: &LlmGenerationRequest,
138    ) -> Result<(serde_json::Value, LlmUsage)> {
139        self.ensure_initialized().await?;
140
141        let engine = self.rag_engine.read().await;
142        let provider = engine
143            .as_ref()
144            .ok_or_else(|| mockforge_foundation::Error::internal("LLM provider not initialized"))?;
145
146        // Build messages
147        let messages = vec![
148            ChatMessage {
149                role: "system".to_string(),
150                content: request.system_prompt.clone(),
151            },
152            ChatMessage {
153                role: "user".to_string(),
154                content: request.user_prompt.clone(),
155            },
156        ];
157
158        // Generate response with usage tracking
159        let (response_text, usage) = provider
160            .generate_chat_with_usage(
161                messages,
162                request.temperature,
163                request.max_tokens,
164                self.resolve_seed(request),
165            )
166            .await?;
167
168        // Try to parse as JSON
169        let json_value = match serde_json::from_str::<serde_json::Value>(&response_text) {
170            Ok(json) => json,
171            Err(_) => {
172                // Try to extract JSON from response
173                if let Some(start) = response_text.find('{') {
174                    if let Some(end) = response_text.rfind('}') {
175                        let json_str = &response_text[start..=end];
176                        if let Ok(json) = serde_json::from_str::<serde_json::Value>(json_str) {
177                            json
178                        } else {
179                            serde_json::json!({
180                                "response": response_text,
181                                "note": "Response was not valid JSON, wrapped in object"
182                            })
183                        }
184                    } else {
185                        serde_json::json!({
186                            "response": response_text,
187                            "note": "Response was not valid JSON, wrapped in object"
188                        })
189                    }
190                } else {
191                    serde_json::json!({
192                        "response": response_text,
193                        "note": "Response was not valid JSON, wrapped in object"
194                    })
195                }
196            }
197        };
198
199        Ok((json_value, usage))
200    }
201
202    /// Get configuration
203    pub fn config(&self) -> &BehaviorModelConfig {
204        &self.config
205    }
206}
207
208/// Chat message for LLM
209#[derive(Debug, Clone)]
210struct ChatMessage {
211    role: String,
212    content: String,
213}
214
215/// LLM usage information
216#[derive(Debug, Clone, Default)]
217pub struct LlmUsage {
218    /// Prompt tokens used
219    pub prompt_tokens: u64,
220    /// Completion tokens used
221    pub completion_tokens: u64,
222    /// Total tokens used
223    pub total_tokens: u64,
224}
225
226impl LlmUsage {
227    /// Create new usage info
228    pub fn new(prompt_tokens: u64, completion_tokens: u64) -> Self {
229        Self {
230            prompt_tokens,
231            completion_tokens,
232            total_tokens: prompt_tokens + completion_tokens,
233        }
234    }
235}
236
237/// LLM provider trait
238#[async_trait::async_trait]
239trait LlmProvider: Send + Sync {
240    /// Generate chat completion
241    async fn generate_chat(
242        &self,
243        messages: Vec<ChatMessage>,
244        temperature: f64,
245        max_tokens: usize,
246        seed: Option<i64>,
247    ) -> Result<String>;
248
249    /// Generate chat completion with usage tracking
250    async fn generate_chat_with_usage(
251        &self,
252        messages: Vec<ChatMessage>,
253        temperature: f64,
254        max_tokens: usize,
255        seed: Option<i64>,
256    ) -> Result<(String, LlmUsage)> {
257        // Default implementation: call generate_chat and estimate tokens
258        let response = self.generate_chat(messages, temperature, max_tokens, seed).await?;
259        // Rough estimation: ~4 characters per token
260        let estimated_tokens = (response.len() as f64 / 4.0) as u64;
261        Ok((response, LlmUsage::new(estimated_tokens, estimated_tokens)))
262    }
263}
264
265/// OpenAI provider implementation
266struct OpenAIProvider {
267    client: reqwest::Client,
268    api_key: String,
269    model: String,
270    endpoint: String,
271}
272
273impl OpenAIProvider {
274    fn new(config: &BehaviorModelConfig) -> Result<Self> {
275        let api_key = config
276            .api_key
277            .clone()
278            .or_else(|| std::env::var("OPENAI_API_KEY").ok())
279            .ok_or_else(|| mockforge_foundation::Error::internal("OpenAI API key not found"))?;
280
281        let endpoint = config
282            .api_endpoint
283            .clone()
284            .unwrap_or_else(|| "https://api.openai.com/v1/chat/completions".to_string());
285
286        Ok(Self {
287            client: reqwest::Client::new(),
288            api_key,
289            model: config.model.clone(),
290            endpoint,
291        })
292    }
293}
294
295#[async_trait::async_trait]
296impl LlmProvider for OpenAIProvider {
297    async fn generate_chat(
298        &self,
299        messages: Vec<ChatMessage>,
300        temperature: f64,
301        max_tokens: usize,
302        seed: Option<i64>,
303    ) -> Result<String> {
304        let mut request_body = serde_json::json!({
305            "model": self.model,
306            "messages": messages.iter().map(|m| {
307                serde_json::json!({
308                    "role": m.role,
309                    "content": m.content
310                })
311            }).collect::<Vec<_>>(),
312            "temperature": temperature,
313            "max_tokens": max_tokens,
314        });
315        if let Some(seed) = seed {
316            request_body["seed"] = serde_json::json!(seed);
317        }
318
319        let response = self
320            .client
321            .post(&self.endpoint)
322            .header("Authorization", format!("Bearer {}", self.api_key))
323            .header("Content-Type", "application/json")
324            .json(&request_body)
325            .send()
326            .await
327            .map_err(|e| {
328                mockforge_foundation::Error::internal(format!("OpenAI API request failed: {}", e))
329            })?;
330
331        if !response.status().is_success() {
332            let error_text = response.text().await.unwrap_or_default();
333            return Err(mockforge_foundation::Error::internal(format!(
334                "OpenAI API error: {}",
335                error_text
336            )));
337        }
338
339        let response_json: serde_json::Value = response.json().await.map_err(|e| {
340            mockforge_foundation::Error::internal(format!("Failed to parse OpenAI response: {}", e))
341        })?;
342
343        // Extract content from response
344        let content = response_json["choices"][0]["message"]["content"]
345            .as_str()
346            .ok_or_else(|| mockforge_foundation::Error::internal("Invalid OpenAI response format"))?
347            .to_string();
348
349        Ok(content)
350    }
351
352    async fn generate_chat_with_usage(
353        &self,
354        messages: Vec<ChatMessage>,
355        temperature: f64,
356        max_tokens: usize,
357        seed: Option<i64>,
358    ) -> Result<(String, LlmUsage)> {
359        let mut request_body = serde_json::json!({
360            "model": self.model,
361            "messages": messages.iter().map(|m| {
362                serde_json::json!({
363                    "role": m.role,
364                    "content": m.content
365                })
366            }).collect::<Vec<_>>(),
367            "temperature": temperature,
368            "max_tokens": max_tokens,
369        });
370        if let Some(seed) = seed {
371            request_body["seed"] = serde_json::json!(seed);
372        }
373
374        let response = self
375            .client
376            .post(&self.endpoint)
377            .header("Authorization", format!("Bearer {}", self.api_key))
378            .header("Content-Type", "application/json")
379            .json(&request_body)
380            .send()
381            .await
382            .map_err(|e| {
383                mockforge_foundation::Error::internal(format!("OpenAI API request failed: {}", e))
384            })?;
385
386        if !response.status().is_success() {
387            let error_text = response.text().await.unwrap_or_default();
388            return Err(mockforge_foundation::Error::internal(format!(
389                "OpenAI API error: {}",
390                error_text
391            )));
392        }
393
394        let response_json: serde_json::Value = response.json().await.map_err(|e| {
395            mockforge_foundation::Error::internal(format!("Failed to parse OpenAI response: {}", e))
396        })?;
397
398        // Extract content from response
399        let content = response_json["choices"][0]["message"]["content"]
400            .as_str()
401            .ok_or_else(|| mockforge_foundation::Error::internal("Invalid OpenAI response format"))?
402            .to_string();
403
404        // Extract usage information
405        let usage = if let Some(usage_obj) = response_json.get("usage") {
406            LlmUsage::new(
407                usage_obj["prompt_tokens"].as_u64().unwrap_or(0),
408                usage_obj["completion_tokens"].as_u64().unwrap_or(0),
409            )
410        } else {
411            // Fallback: estimate tokens
412            let estimated = (content.len() as f64 / 4.0) as u64;
413            LlmUsage::new(estimated, estimated)
414        };
415
416        Ok((content, usage))
417    }
418}
419
420/// Ollama provider implementation
421struct OllamaProvider {
422    client: reqwest::Client,
423    model: String,
424    endpoint: String,
425}
426
427impl OllamaProvider {
428    fn new(config: &BehaviorModelConfig) -> Result<Self> {
429        let endpoint = config
430            .api_endpoint
431            .clone()
432            .unwrap_or_else(|| "http://localhost:11434/api/chat".to_string());
433
434        Ok(Self {
435            client: reqwest::Client::new(),
436            model: config.model.clone(),
437            endpoint,
438        })
439    }
440}
441
442#[async_trait::async_trait]
443impl LlmProvider for OllamaProvider {
444    async fn generate_chat(
445        &self,
446        messages: Vec<ChatMessage>,
447        temperature: f64,
448        max_tokens: usize,
449        seed: Option<i64>,
450    ) -> Result<String> {
451        let mut request_body = serde_json::json!({
452            "model": self.model,
453            "messages": messages.iter().map(|m| {
454                serde_json::json!({
455                    "role": m.role,
456                    "content": m.content
457                })
458            }).collect::<Vec<_>>(),
459            "options": {
460                "temperature": temperature,
461                "num_predict": max_tokens,
462            },
463            "stream": false,
464        });
465        // Ollama takes the seed inside `options` (#852).
466        if let Some(seed) = seed {
467            request_body["options"]["seed"] = serde_json::json!(seed);
468        }
469        let response = self
470            .client
471            .post(&self.endpoint)
472            .header("Content-Type", "application/json")
473            .json(&request_body)
474            .send()
475            .await
476            .map_err(|e| {
477                mockforge_foundation::Error::internal(format!("Ollama API request failed: {}", e))
478            })?;
479
480        if !response.status().is_success() {
481            let error_text = response.text().await.unwrap_or_default();
482            return Err(mockforge_foundation::Error::internal(format!(
483                "Ollama API error: {}",
484                error_text
485            )));
486        }
487
488        let response_json: serde_json::Value = response.json().await.map_err(|e| {
489            mockforge_foundation::Error::internal(format!("Failed to parse Ollama response: {}", e))
490        })?;
491
492        // Extract content from response
493        let content = response_json["message"]["content"]
494            .as_str()
495            .ok_or_else(|| mockforge_foundation::Error::internal("Invalid Ollama response format"))?
496            .to_string();
497
498        Ok(content)
499    }
500}
501
502/// Anthropic provider implementation
503struct AnthropicProvider {
504    client: reqwest::Client,
505    api_key: String,
506    model: String,
507    endpoint: String,
508}
509
510impl AnthropicProvider {
511    fn new(config: &BehaviorModelConfig) -> Result<Self> {
512        let api_key = config
513            .api_key
514            .clone()
515            .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
516            .ok_or_else(|| mockforge_foundation::Error::internal("Anthropic API key not found"))?;
517
518        let endpoint = config
519            .api_endpoint
520            .clone()
521            .unwrap_or_else(|| "https://api.anthropic.com/v1/messages".to_string());
522
523        Ok(Self {
524            client: reqwest::Client::new(),
525            api_key,
526            model: config.model.clone(),
527            endpoint,
528        })
529    }
530}
531
532#[async_trait::async_trait]
533impl LlmProvider for AnthropicProvider {
534    async fn generate_chat(
535        &self,
536        messages: Vec<ChatMessage>,
537        temperature: f64,
538        max_tokens: usize,
539        // Anthropic's Messages API has no sampling-seed parameter; ignored.
540        _seed: Option<i64>,
541    ) -> Result<String> {
542        // Separate system message from other messages
543        let system_message =
544            messages.iter().find(|m| m.role == "system").map(|m| m.content.clone());
545
546        let chat_messages: Vec<_> = messages
547            .iter()
548            .filter(|m| m.role != "system")
549            .map(|m| {
550                serde_json::json!({
551                    "role": m.role,
552                    "content": m.content
553                })
554            })
555            .collect();
556
557        let mut request_body = serde_json::json!({
558            "model": self.model,
559            "messages": chat_messages,
560            "temperature": temperature,
561            "max_tokens": max_tokens,
562        });
563
564        if let Some(system) = system_message {
565            request_body["system"] = serde_json::Value::String(system);
566        }
567
568        let response = self
569            .client
570            .post(&self.endpoint)
571            .header("x-api-key", &self.api_key)
572            .header("anthropic-version", "2023-06-01")
573            .header("Content-Type", "application/json")
574            .json(&request_body)
575            .send()
576            .await
577            .map_err(|e| {
578                mockforge_foundation::Error::internal(format!(
579                    "Anthropic API request failed: {}",
580                    e
581                ))
582            })?;
583
584        if !response.status().is_success() {
585            let error_text = response.text().await.unwrap_or_default();
586            return Err(mockforge_foundation::Error::internal(format!(
587                "Anthropic API error: {}",
588                error_text
589            )));
590        }
591
592        let response_json: serde_json::Value = response.json().await.map_err(|e| {
593            mockforge_foundation::Error::internal(format!(
594                "Failed to parse Anthropic response: {}",
595                e
596            ))
597        })?;
598
599        // Extract content from response
600        let content = response_json["content"][0]["text"]
601            .as_str()
602            .ok_or_else(|| {
603                mockforge_foundation::Error::internal("Invalid Anthropic response format")
604            })?
605            .to_string();
606
607        Ok(content)
608    }
609}
610
611/// OpenAI-compatible provider (generic)
612struct OpenAICompatibleProvider {
613    client: reqwest::Client,
614    api_key: Option<String>,
615    model: String,
616    endpoint: String,
617}
618
619impl OpenAICompatibleProvider {
620    fn new(config: &BehaviorModelConfig) -> Result<Self> {
621        let endpoint = config.api_endpoint.clone().ok_or_else(|| {
622            mockforge_foundation::Error::internal(
623                "API endpoint required for OpenAI-compatible provider",
624            )
625        })?;
626
627        Ok(Self {
628            client: reqwest::Client::new(),
629            api_key: config.api_key.clone(),
630            model: config.model.clone(),
631            endpoint,
632        })
633    }
634}
635
636#[async_trait::async_trait]
637impl LlmProvider for OpenAICompatibleProvider {
638    async fn generate_chat(
639        &self,
640        messages: Vec<ChatMessage>,
641        temperature: f64,
642        max_tokens: usize,
643        seed: Option<i64>,
644    ) -> Result<String> {
645        let mut request_body = serde_json::json!({
646            "model": self.model,
647            "messages": messages.iter().map(|m| {
648                serde_json::json!({
649                    "role": m.role,
650                    "content": m.content
651                })
652            }).collect::<Vec<_>>(),
653            "temperature": temperature,
654            "max_tokens": max_tokens,
655        });
656        // Most OpenAI-compatible servers honour the OpenAI `seed` field
657        // (#852); those that don't ignore it.
658        if let Some(seed) = seed {
659            request_body["seed"] = serde_json::json!(seed);
660        }
661        let mut request =
662            self.client.post(&self.endpoint).header("Content-Type", "application/json");
663
664        if let Some(api_key) = &self.api_key {
665            request = request.header("Authorization", format!("Bearer {}", api_key));
666        }
667
668        let response = request.json(&request_body).send().await.map_err(|e| {
669            mockforge_foundation::Error::internal(format!("API request failed: {}", e))
670        })?;
671
672        if !response.status().is_success() {
673            let error_text = response.text().await.unwrap_or_default();
674            return Err(mockforge_foundation::Error::internal(format!(
675                "API error: {}",
676                error_text
677            )));
678        }
679
680        let response_json: serde_json::Value = response.json().await.map_err(|e| {
681            mockforge_foundation::Error::internal(format!("Failed to parse API response: {}", e))
682        })?;
683
684        // Extract content (try both OpenAI and Ollama formats)
685        let content = response_json["choices"][0]["message"]["content"]
686            .as_str()
687            .or_else(|| response_json["message"]["content"].as_str())
688            .ok_or_else(|| mockforge_foundation::Error::internal("Invalid API response format"))?
689            .to_string();
690
691        Ok(content)
692    }
693}
694
695#[cfg(test)]
696mod tests {
697    use super::*;
698
699    #[test]
700    fn test_llm_client_creation() {
701        let config = BehaviorModelConfig::default();
702        let client = LlmClient::new(config);
703        assert_eq!(client.config().llm_provider, "openai");
704    }
705}
mockforge_intelligence/intelligent_behavior/llm_client.rs

mockforge_intelligence/intelligent_behavior/
llm_client.rs
