orchestral_runtime/planner/
factory.rs

1//! LLM client factory for building clients from backend configuration.
2
3use std::str::FromStr;
4use std::sync::Arc;
5use std::time::Duration;
6
7use async_trait::async_trait;
8use futures_util::StreamExt;
9use llm_sdk::builder::{FunctionBuilder, LLMBackend, LLMBuilder};
10use llm_sdk::chat::{ChatMessage, ToolChoice};
11use thiserror::Error;
12
13use orchestral_core::config::BackendSpec;
14
15use super::llm::{
16    LlmClient, LlmError, LlmRequest, LlmResponse, StreamChunkCallback, ToolDefinition,
17};
18
19/// Runtime invocation config for an LLM call chain.
20#[derive(Debug, Clone)]
21pub struct LlmInvocationConfig {
22    pub model: String,
23    pub temperature: f32,
24    pub max_tokens: u32,
25    pub normalize_response: bool,
26}
27
28impl Default for LlmInvocationConfig {
29    fn default() -> Self {
30        Self {
31            model: "anthropic/claude-sonnet-4.5".to_string(),
32            temperature: 0.2,
33            max_tokens: 4096,
34            normalize_response: true,
35        }
36    }
37}
38
39/// Errors that can occur when building an LLM client.
40#[derive(Debug, Error)]
41pub enum LlmBuildError {
42    #[error("unknown backend kind: {0}")]
43    UnknownKind(String),
44    #[error("missing API key for backend")]
45    MissingApiKey,
46    #[error("environment variable '{0}' not found")]
47    EnvNotFound(String),
48}
49
50/// Factory trait for building LLM clients.
51pub trait LlmClientFactory: Send + Sync {
52    /// Build an LLM client from backend + invocation config.
53    fn build(
54        &self,
55        backend: &BackendSpec,
56        invocation: &LlmInvocationConfig,
57    ) -> Result<Arc<dyn LlmClient>, LlmBuildError>;
58}
59
60/// Default factory implementation using `graniet/llm`.
61pub struct DefaultLlmClientFactory;
62
63impl DefaultLlmClientFactory {
64    pub fn new() -> Self {
65        Self
66    }
67}
68
69impl Default for DefaultLlmClientFactory {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl LlmClientFactory for DefaultLlmClientFactory {
76    fn build(
77        &self,
78        backend: &BackendSpec,
79        invocation: &LlmInvocationConfig,
80    ) -> Result<Arc<dyn LlmClient>, LlmBuildError> {
81        build_client_from_backend(backend, invocation)
82    }
83}
84
85/// Build an LLM client from backend spec and invocation defaults.
86pub fn build_client_from_backend(
87    backend: &BackendSpec,
88    invocation: &LlmInvocationConfig,
89) -> Result<Arc<dyn LlmClient>, LlmBuildError> {
90    let backend_kind = parse_backend(&backend.kind)?;
91    let api_key = resolve_api_key(backend)?;
92    let client = GranietLlmClient {
93        backend: backend_kind,
94        api_key: Some(api_key),
95        base_url: backend.endpoint.clone(),
96        model: invocation.model.clone(),
97        temperature: invocation.temperature,
98        max_tokens: backend
99            .get_config::<u32>("max_tokens")
100            .unwrap_or(invocation.max_tokens),
101        normalize_response: invocation.normalize_response,
102        timeout_secs: backend.get_config::<u64>("timeout_secs").unwrap_or(60),
103    };
104    Ok(Arc::new(client))
105}
106
107fn resolve_api_key(spec: &BackendSpec) -> Result<String, LlmBuildError> {
108    let candidates = api_key_env_candidates(spec);
109    let first = candidates
110        .first()
111        .cloned()
112        .ok_or(LlmBuildError::MissingApiKey)?;
113    for env_name in candidates {
114        if let Ok(value) = std::env::var(&env_name) {
115            if !value.trim().is_empty() {
116                return Ok(value);
117            }
118        }
119    }
120    Err(LlmBuildError::EnvNotFound(first))
121}
122
123fn api_key_env_candidates(spec: &BackendSpec) -> Vec<String> {
124    let mut candidates = Vec::new();
125    if let Some(explicit) = spec.api_key_env.as_ref() {
126        candidates.push(explicit.clone());
127    }
128    for fallback in default_api_key_envs_for_kind(&spec.kind) {
129        if !candidates.iter().any(|existing| existing == fallback) {
130            candidates.push(fallback.to_string());
131        }
132    }
133    candidates
134}
135
136fn default_api_key_envs_for_kind(kind: &str) -> &'static [&'static str] {
137    match kind.trim().to_ascii_lowercase().as_str() {
138        "openai" => &["OPENAI_API_KEY"],
139        "google" | "gemini" => &["GOOGLE_API_KEY", "GEMINI_API_KEY"],
140        "anthropic" | "claude" => &["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
141        "deepseek" => &["DEEPSEEK_API_KEY"],
142        "groq" => &["GROQ_API_KEY"],
143        "xai" => &["XAI_API_KEY"],
144        "mistral" => &["MISTRAL_API_KEY"],
145        "cohere" => &["COHERE_API_KEY"],
146        "openrouter" => &["OPENROUTER_API_KEY"],
147        _ => &[],
148    }
149}
150
151fn parse_backend(kind: &str) -> Result<LLMBackend, LlmBuildError> {
152    let normalized = kind.to_lowercase();
153    let mapped = match normalized.as_str() {
154        "gemini" => "google",
155        "claude" => "anthropic",
156        other => other,
157    };
158    let allowed = matches!(
159        mapped,
160        "openai"
161            | "google"
162            | "anthropic"
163            | "deepseek"
164            | "groq"
165            | "xai"
166            | "mistral"
167            | "cohere"
168            | "openrouter"
169            | "ollama"
170    );
171    if !allowed {
172        return Err(LlmBuildError::UnknownKind(kind.to_string()));
173    }
174    LLMBackend::from_str(mapped).map_err(|_| LlmBuildError::UnknownKind(kind.to_string()))
175}
176
177struct GranietLlmClient {
178    backend: LLMBackend,
179    api_key: Option<String>,
180    base_url: Option<String>,
181    model: String,
182    temperature: f32,
183    max_tokens: u32,
184    normalize_response: bool,
185    timeout_secs: u64,
186}
187
188#[async_trait]
189impl LlmClient for GranietLlmClient {
190    async fn complete(&self, request: LlmRequest) -> Result<String, LlmError> {
191        let prompt = if request.system.trim().is_empty() {
192            request.user
193        } else {
194            format!("System:\n{}\n\nUser:\n{}", request.system, request.user)
195        };
196
197        let model = if request.model.trim().is_empty() {
198            self.model.clone()
199        } else {
200            request.model
201        };
202        let temperature = if request.temperature <= 0.0 {
203            self.temperature
204        } else {
205            request.temperature
206        };
207
208        let mut builder = LLMBuilder::new()
209            .backend(self.backend.clone())
210            .model(model)
211            .temperature(temperature)
212            .max_tokens(self.max_tokens)
213            .normalize_response(self.normalize_response);
214        if let Some(endpoint) = &self.base_url {
215            builder = builder.base_url(endpoint.clone());
216        }
217        if let Some(api_key) = &self.api_key {
218            builder = builder.api_key(api_key.clone());
219        }
220
221        let llm = builder
222            .build()
223            .map_err(|e| LlmError::Http(format!("llm builder error: {}", e)))?;
224
225        let messages = vec![ChatMessage::user().content(prompt).build()];
226        let response =
227            tokio::time::timeout(Duration::from_secs(self.timeout_secs), llm.chat(&messages))
228                .await
229                .map_err(|_| {
230                    LlmError::Http(format!("llm chat timeout after {}s", self.timeout_secs))
231                })?
232                .map_err(|e| LlmError::Http(format!("llm chat error: {}", e)))?;
233
234        response
235            .text()
236            .map(|s| s.to_string())
237            .ok_or_else(|| LlmError::Response("llm response had no text".to_string()))
238    }
239
240    async fn complete_with_tools(
241        &self,
242        request: LlmRequest,
243        tools: &[ToolDefinition],
244    ) -> Result<LlmResponse, LlmError> {
245        let prompt = if request.system.trim().is_empty() {
246            request.user
247        } else {
248            format!("System:\n{}\n\nUser:\n{}", request.system, request.user)
249        };
250
251        let model = if request.model.trim().is_empty() {
252            self.model.clone()
253        } else {
254            request.model
255        };
256        let temperature = if request.temperature <= 0.0 {
257            self.temperature
258        } else {
259            request.temperature
260        };
261
262        let mut builder = LLMBuilder::new()
263            .backend(self.backend.clone())
264            .model(model)
265            .temperature(temperature)
266            .max_tokens(self.max_tokens)
267            .normalize_response(self.normalize_response);
268        for t in tools {
269            builder = builder.function(
270                FunctionBuilder::new(&t.name)
271                    .description(&t.description)
272                    .json_schema(t.parameters.clone()),
273            );
274        }
275        builder = builder.tool_choice(ToolChoice::Any);
276        if let Some(endpoint) = &self.base_url {
277            builder = builder.base_url(endpoint.clone());
278        }
279        if let Some(api_key) = &self.api_key {
280            builder = builder.api_key(api_key.clone());
281        }
282
283        let llm = builder
284            .build()
285            .map_err(|e| LlmError::Http(format!("llm builder error: {}", e)))?;
286
287        let messages = vec![ChatMessage::user().content(prompt).build()];
288        let response =
289            tokio::time::timeout(Duration::from_secs(self.timeout_secs), llm.chat(&messages))
290                .await
291                .map_err(|_| {
292                    LlmError::Http(format!("llm chat timeout after {}s", self.timeout_secs))
293                })?
294                .map_err(|e| LlmError::Http(format!("llm chat_with_tools error: {}", e)))?;
295
296        if let Some(tool_calls) = response.tool_calls() {
297            if let Some(tc) = tool_calls.into_iter().next() {
298                let arguments: serde_json::Value = serde_json::from_str(&tc.function.arguments)
299                    .map_err(|e| {
300                        LlmError::Serialization(format!(
301                            "failed to parse tool call arguments: {}",
302                            e
303                        ))
304                    })?;
305                return Ok(LlmResponse::ToolCall {
306                    id: tc.id.clone(),
307                    name: tc.function.name.clone(),
308                    arguments,
309                });
310            }
311        }
312
313        let text = response.text().unwrap_or_default();
314        Ok(LlmResponse::Text(text))
315    }
316
317    async fn complete_stream(
318        &self,
319        request: LlmRequest,
320        on_chunk: StreamChunkCallback,
321    ) -> Result<String, LlmError> {
322        let prompt = if request.system.trim().is_empty() {
323            request.user
324        } else {
325            format!("System:\n{}\n\nUser:\n{}", request.system, request.user)
326        };
327        let model = if request.model.trim().is_empty() {
328            self.model.clone()
329        } else {
330            request.model
331        };
332        let temperature = if request.temperature <= 0.0 {
333            self.temperature
334        } else {
335            request.temperature
336        };
337
338        let mut builder = LLMBuilder::new()
339            .backend(self.backend.clone())
340            .model(model)
341            .temperature(temperature)
342            .max_tokens(self.max_tokens)
343            .normalize_response(self.normalize_response);
344        if let Some(endpoint) = &self.base_url {
345            builder = builder.base_url(endpoint.clone());
346        }
347        if let Some(api_key) = &self.api_key {
348            builder = builder.api_key(api_key.clone());
349        }
350        let llm = builder
351            .build()
352            .map_err(|e| LlmError::Http(format!("llm builder error: {}", e)))?;
353        let messages = vec![ChatMessage::user().content(prompt).build()];
354
355        let mut stream = tokio::time::timeout(
356            Duration::from_secs(self.timeout_secs),
357            llm.chat_stream(&messages),
358        )
359        .await
360        .map_err(|_| {
361            LlmError::Http(format!(
362                "llm stream setup timeout after {}s",
363                self.timeout_secs
364            ))
365        })?
366        .map_err(|e| LlmError::Http(format!("llm chat_stream error: {}", e)))?;
367
368        let mut full = String::new();
369        while let Some(item) =
370            tokio::time::timeout(Duration::from_secs(self.timeout_secs), stream.next())
371                .await
372                .map_err(|_| {
373                    LlmError::Http(format!(
374                        "llm stream timeout after {}s while reading chunk",
375                        self.timeout_secs
376                    ))
377                })?
378        {
379            let chunk =
380                item.map_err(|e| LlmError::Http(format!("llm stream chunk error: {}", e)))?;
381            if chunk.is_empty() {
382                continue;
383            }
384            full.push_str(&chunk);
385            on_chunk(chunk);
386        }
387        if full.is_empty() {
388            return Err(LlmError::Response(
389                "llm stream produced no text chunks".to_string(),
390            ));
391        }
392        Ok(full)
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399    use serde_json::json;
400
401    fn make_backend(kind: &str) -> BackendSpec {
402        BackendSpec {
403            name: "test".to_string(),
404            kind: kind.to_string(),
405            endpoint: None,
406            api_key_env: Some("TEST_API_KEY".to_string()),
407            config: json!({}),
408        }
409    }
410
411    #[test]
412    fn test_unknown_kind() {
413        let backend = make_backend("not-a-real-backend-kind");
414        let invocation = LlmInvocationConfig::default();
415        std::env::set_var("TEST_API_KEY", "dummy");
416        let result = build_client_from_backend(&backend, &invocation);
417        std::env::remove_var("TEST_API_KEY");
418        assert!(matches!(result, Err(LlmBuildError::UnknownKind(_))));
419    }
420
421    #[test]
422    fn test_missing_env_var() {
423        let backend = make_backend("openai");
424        let invocation = LlmInvocationConfig::default();
425        std::env::remove_var("TEST_API_KEY");
426        std::env::remove_var("OPENAI_API_KEY");
427        let result = build_client_from_backend(&backend, &invocation);
428        assert!(matches!(result, Err(LlmBuildError::EnvNotFound(_))));
429    }
430
431    #[test]
432    fn test_google_backend_accepts_google_api_key_alias() {
433        let backend = BackendSpec {
434            name: "google".to_string(),
435            kind: "google".to_string(),
436            endpoint: None,
437            api_key_env: None,
438            config: json!({}),
439        };
440        std::env::remove_var("GEMINI_API_KEY");
441        std::env::set_var("GOOGLE_API_KEY", "google-key");
442        let resolved = resolve_api_key(&backend).expect("google api key");
443        std::env::remove_var("GOOGLE_API_KEY");
444        assert_eq!(resolved, "google-key");
445    }
446
447    #[test]
448    fn test_claude_backend_accepts_claude_api_key_alias() {
449        let backend = BackendSpec {
450            name: "claude".to_string(),
451            kind: "claude".to_string(),
452            endpoint: None,
453            api_key_env: None,
454            config: json!({}),
455        };
456        std::env::remove_var("ANTHROPIC_API_KEY");
457        std::env::set_var("CLAUDE_API_KEY", "claude-key");
458        let resolved = resolve_api_key(&backend).expect("claude api key");
459        std::env::remove_var("CLAUDE_API_KEY");
460        assert_eq!(resolved, "claude-key");
461        assert!(parse_backend("claude").is_ok());
462    }
463}
orchestral_runtime/planner/factory.rs

orchestral_runtime/planner/
factory.rs
