ares_llm/
client.rs

1use crate::config::{ModelConfig, ProviderConfig};
2use ares_types::types::{AppError, Result, ToolCall, ToolDefinition};
3use async_trait::async_trait;
4use genai::adapter::AdapterKind;
5use std::collections::HashMap;
6
7/// Azure AI Foundry env/header helpers (no LLMClient). Inlined after
8/// `azure.rs` was removed from this crate.
9pub(crate) const AZURE_API_KEY_ENV: &str = "AZURE_FOUNDRY_API_KEY";
10pub(crate) const AZURE_BASE_URL_ENV: &str = "AZURE_FOUNDRY_BASE_URL";
11pub(crate) const AZURE_MODEL_ENV: &str = "AZURE_FOUNDRY_MODEL";
12pub(crate) const AZURE_DEFAULT_MODEL: &str = "DeepSeek-V4-Flash";
13const AZURE_MODEL_PREFIX: &str = "azure/";
14
15pub(crate) fn azure_strip_model_prefix(model: &str) -> &str {
16    let trimmed = model.trim();
17    trimmed
18        .strip_prefix(AZURE_MODEL_PREFIX)
19        .map(str::trim)
20        .filter(|model| !model.is_empty())
21        .unwrap_or(trimmed)
22}
23
24pub(crate) fn azure_normalize_base_url(api_base: &str) -> String {
25    api_base.trim().trim_end_matches('/').to_string()
26}
27
28pub(crate) fn azure_foundry_headers(api_key: &str) -> HashMap<String, String> {
29    let mut headers = HashMap::with_capacity(2);
30    headers.insert("api-key".to_string(), api_key.to_string());
31    headers.insert("Authorization".to_string(), format!("Bearer {api_key}"));
32    headers
33}
34
35/// Provider-neutral prompt cache policy mapped onto genai cache control.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
37#[serde(rename_all = "snake_case")]
38pub enum CacheControl {
39    /// Default ephemeral cache.
40    Ephemeral,
41    /// Explicit 5-minute TTL.
42    Ephemeral5m,
43    /// Extended 24-hour TTL.
44    Ephemeral24h,
45}
46
47/// Optional generation hints set on a client between calls.
48///
49/// Hints are OPT-IN: [`LLMClient::supports_hints`] defaults to `false` and
50/// [`LLMClient::set_hints`] defaults to a no-op, so every existing provider
51/// implementation keeps compiling and behaving exactly as before. A provider
52/// adopts hints by overriding both methods (and honoring the stored hints in
53/// its request building).
54///
55/// # Set-on-client semantics
56///
57/// `generate_with_system`'s signature is fixed and widely called, so hints
58/// ride ON THE CLIENT instead of through per-call parameters: they apply to
59/// all SUBSEQUENT generate calls until replaced by another `set_hints`
60/// (clear with `GenerationHints::default()`).
61///
62/// # Thread safety
63///
64/// Implementations that adopt hints MUST use interior mutability (for
65/// example `std::sync::RwLock<GenerationHints>`) because the trait methods
66/// take `&self`. Readers snapshot the hints at the start of each call.
67#[derive(Debug, Clone, Default, PartialEq, Eq)]
68pub struct GenerationHints {
69    /// Ask the provider for a JSON object response.
70    pub json_mode: bool,
71    /// Ask a reasoning-capable model to skip visible reasoning output.
72    pub suppress_reasoning: bool,
73    /// Advisory maximum number of output tokens (`None` = provider default).
74    pub max_tokens: Option<u32>,
75    /// Optional constrained-output grammar in GBNF/EBNF-style syntax
76    /// (`None` = unconstrained). Providers honor it only where a native
77    /// mechanism exists; unsupported backends silently ignore it without
78    /// erroring.
79    pub guided_grammar: Option<String>,
80    /// Reasoning effort keyword (`low|medium|high|zero`).
81    pub reasoning_effort: Option<String>,
82    /// OpenAI prompt cache key.
83    pub prompt_cache_key: Option<String>,
84    /// Request-level cache control.
85    pub cache_control: Option<CacheControl>,
86    /// Previous Responses API id for stateful continuation.
87    pub previous_response_id: Option<String>,
88    /// Whether the provider should store the response.
89    pub store: Option<bool>,
90    /// Attach the provider built-in web search tool (`provider_web_search`).
91    pub web_search: bool,
92}
93
94/// Generic LLM client trait for provider abstraction
95#[async_trait]
96pub trait LLMClient: Send + Sync {
97    /// Generate a completion from a prompt
98    async fn generate(&self, prompt: &str) -> Result<String>;
99
100    /// Generate with system prompt
101    async fn generate_with_system(&self, system: &str, prompt: &str) -> Result<String>;
102
103    /// Generate with conversation history, returning full response with token usage
104    async fn generate_with_history(
105        &self,
106        messages: &[(String, String)], // (role, content) pairs
107    ) -> Result<LLMResponse>;
108
109    /// Generate with tool calling support
110    async fn generate_with_tools(
111        &self,
112        prompt: &str,
113        tools: &[ToolDefinition],
114    ) -> Result<LLMResponse>;
115
116    /// Generate with conversation history AND tool definitions.
117    async fn generate_with_tools_and_history(
118        &self,
119        messages: &[crate::coordinator::ConversationMessage],
120        tools: &[ToolDefinition],
121    ) -> Result<LLMResponse>;
122
123    /// Stream a completion
124    async fn stream(
125        &self,
126        prompt: &str,
127    ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
128
129    /// Stream a completion with system prompt
130    async fn stream_with_system(
131        &self,
132        system: &str,
133        prompt: &str,
134    ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
135
136    /// Stream a completion with conversation history
137    async fn stream_with_history(
138        &self,
139        messages: &[(String, String)], // (role, content) pairs
140    ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
141
142    /// Get the model name/identifier
143    fn model_name(&self) -> &str;
144
145    /// Whether this client honors [`GenerationHints`] set via
146    /// [`LLMClient::set_hints`]. Defaults to `false`; hint-aware providers
147    /// override this together with `set_hints`.
148    fn supports_hints(&self) -> bool {
149        false
150    }
151
152    /// Store generation hints applying to SUBSEQUENT generate calls, until
153    /// replaced (clear with `GenerationHints::default()`). Default impl is a
154    /// no-op so unmodified providers keep compiling unchanged.
155    fn set_hints(&self, _hints: GenerationHints) {}
156
157    /// Embed one or more input strings. Default: not supported.
158    async fn embed(&self, _inputs: &[String]) -> Result<Vec<Vec<f32>>> {
159        Err(AppError::FeatureDisabled(
160            "embeddings not supported by this client".into(),
161        ))
162    }
163
164    /// Whether this client can send image/file parts.
165    fn supports_vision(&self) -> bool {
166        false
167    }
168
169    /// Whether this client can attach the provider built-in web search tool.
170    fn supports_provider_web_search(&self) -> bool {
171        false
172    }
173}
174
175/// Token usage statistics from an LLM generation call
176#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
177pub struct TokenUsage {
178    /// Number of tokens in the prompt/input
179    pub prompt_tokens: u32,
180    /// Number of tokens in the completion/output
181    pub completion_tokens: u32,
182    /// Total tokens used (prompt + completion)
183    pub total_tokens: u32,
184    /// Tokens served from the provider-side prompt cache, when the provider
185    /// reports cache hits (`None` when unknown or not reported). Always `0`
186    /// or more; cache hits are a subset of `prompt_tokens`.
187    #[serde(default)]
188    pub cached_tokens: Option<i64>,
189}
190
191impl TokenUsage {
192    /// Create a new TokenUsage with the given values
193    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
194        Self {
195            prompt_tokens,
196            completion_tokens,
197            total_tokens: prompt_tokens + completion_tokens,
198            cached_tokens: None,
199        }
200    }
201}
202
203/// Response from an LLM generation call
204#[derive(Debug, Clone)]
205pub struct LLMResponse {
206    /// The generated text content
207    pub content: String,
208    /// Any tool calls the model wants to make
209    pub tool_calls: Vec<ToolCall>,
210    /// Reason the generation finished (e.g., "stop", "tool_calls", "length")
211    pub finish_reason: String,
212    /// Token usage statistics (if provided by the model)
213    pub usage: Option<TokenUsage>,
214    /// Reasoning/thinking content when the model reports it.
215    pub reasoning_content: Option<String>,
216    /// Provider response id for stateful continuation (Responses API).
217    pub response_id: Option<String>,
218}
219
220/// Model inference parameters
221#[derive(Debug, Clone, PartialEq, Default)]
222pub struct ModelParams {
223    /// Sampling temperature (0.0 = deterministic, 1.0+ = creative)
224    pub temperature: Option<f32>,
225    /// Maximum tokens to generate
226    pub max_tokens: Option<u32>,
227    /// Nucleus sampling parameter
228    pub top_p: Option<f32>,
229    /// Frequency penalty (-2.0 to 2.0)
230    pub frequency_penalty: Option<f32>,
231    /// Presence penalty (-2.0 to 2.0)
232    pub presence_penalty: Option<f32>,
233}
234
235impl ModelParams {
236    /// Create params from a ModelConfig
237    pub fn from_model_config(config: &ModelConfig) -> Self {
238        Self {
239            temperature: Some(config.temperature),
240            max_tokens: Some(config.max_tokens),
241            top_p: None,
242            frequency_penalty: None,
243            presence_penalty: None,
244        }
245    }
246}
247
248/// Resolved genai HTTP provider (kind + credentials + endpoint).
249#[derive(Debug, Clone)]
250pub struct GenaiProvider {
251    /// genai adapter kind used for every call.
252    pub kind: AdapterKind,
253    /// API key (None for unauthenticated local adapters).
254    pub api_key: Option<String>,
255    /// Override endpoint; None uses the adapter default.
256    pub endpoint: Option<String>,
257    /// Model identifier.
258    pub model: String,
259    /// Sampling parameters.
260    pub params: ModelParams,
261    /// Extra HTTP headers (Azure Foundry, runtime providers).
262    pub headers: HashMap<String, String>,
263    /// AWS region (Bedrock API).
264    pub region: Option<String>,
265    /// GCP project (Vertex).
266    pub vertex_project: Option<String>,
267    /// Vertex location.
268    pub vertex_location: Option<String>,
269    /// Custom adapter index (`GENAI_{n}_*`).
270    pub custom_index: Option<u8>,
271}
272
273impl GenaiProvider {
274    fn openai(
275        api_key: String,
276        endpoint: String,
277        model: String,
278        params: ModelParams,
279        headers: HashMap<String, String>,
280    ) -> Self {
281        Self {
282            kind: AdapterKind::OpenAI,
283            api_key: Some(api_key),
284            endpoint: Some(endpoint),
285            model,
286            params,
287            headers,
288            region: None,
289            vertex_project: None,
290            vertex_location: None,
291            custom_index: None,
292        }
293    }
294}
295
296/// LLM Provider configuration
297#[derive(Debug, Clone)]
298#[non_exhaustive]
299#[allow(clippy::large_enum_variant)] // GenaiProvider owns the genai Client
300pub enum Provider {
301    /// Any HTTP provider routed through genai.
302    Genai(GenaiProvider),
303    /// Local GGUF inference via llama.cpp.
304    #[cfg(feature = "llamacpp")]
305    LlamaCpp {
306        /// Path to a GGUF model file.
307        model_path: String,
308        /// Model inference parameters.
309        params: ModelParams,
310    },
311    /// In-memory stub for unit tests (no network I/O).
312    #[cfg(test)]
313    TestStub {
314        /// Model label returned by [`LLMClient::model_name`].
315        model: String,
316    },
317}
318
319impl Provider {
320    /// Create an LLM client from this provider configuration
321    pub async fn create_client(&self) -> Result<Box<dyn LLMClient>> {
322        match self {
323            Provider::Genai(provider) => Ok(Box::new(crate::genai_client::GenaiClient::new(
324                provider.clone(),
325            )?)),
326            #[cfg(feature = "llamacpp")]
327            Provider::LlamaCpp { model_path, params } => Ok(Box::new(
328                crate::llamacpp::LlamaCppClient::with_params(model_path.clone(), params.clone())?,
329            )),
330            #[cfg(test)]
331            Provider::TestStub { model } => {
332                Ok(Box::new(test_support::MockLLMClient::new(model.clone())))
333            }
334        }
335    }
336
337    /// Create a provider from environment variables.
338    ///
339    /// Priority: OPENAI_API_KEY, NVIDIA_API_KEY, AZURE_FOUNDRY_API_KEY,
340    /// AWS_BEARER_TOKEN_BEDROCK, ANTHROPIC_API_KEY, GEMINI_API_KEY, else
341    /// Ollama localhost.
342    pub fn from_env() -> Result<Self> {
343        if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
344            if !api_key.is_empty() {
345                let api_base = std::env::var("OPENAI_API_BASE")
346                    .unwrap_or_else(|_| "https://api.openai.com/v1".into());
347                let model = std::env::var("OPENAI_MODEL").unwrap_or_else(|_| "gpt-4".into());
348                return Ok(Provider::Genai(GenaiProvider::openai(
349                    api_key,
350                    api_base,
351                    model,
352                    ModelParams::default(),
353                    HashMap::new(),
354                )));
355            }
356        }
357
358        if let Ok(api_key) = std::env::var("NVIDIA_API_KEY") {
359            if !api_key.is_empty() {
360                return Ok(Provider::Genai(GenaiProvider::openai(
361                    api_key,
362                    "https://integrate.api.nvidia.com/v1".into(),
363                    "nvidia/nemotron-3-ultra-550b-a55b".into(),
364                    ModelParams::default(),
365                    HashMap::new(),
366                )));
367            }
368        }
369
370        if let Ok(api_key) = std::env::var(AZURE_API_KEY_ENV) {
371            if !api_key.is_empty() {
372                let api_base = std::env::var(AZURE_BASE_URL_ENV).map_err(|_| {
373                    AppError::Configuration(format!(
374                        "{} must be set when {} is configured",
375                        AZURE_BASE_URL_ENV, AZURE_API_KEY_ENV
376                    ))
377                })?;
378                let model = std::env::var(AZURE_MODEL_ENV)
379                    .unwrap_or_else(|_| AZURE_DEFAULT_MODEL.to_string());
380                return Ok(Provider::Genai(GenaiProvider::openai(
381                    api_key.clone(),
382                    azure_normalize_base_url(&api_base),
383                    azure_strip_model_prefix(&model).to_string(),
384                    ModelParams::default(),
385                    azure_foundry_headers(&api_key),
386                )));
387            }
388        }
389
390        if let Ok(api_key) = std::env::var("AWS_BEARER_TOKEN_BEDROCK") {
391            if !api_key.is_empty() {
392                let region = std::env::var("AWS_REGION").map_err(|_| {
393                    AppError::Configuration(
394                        "AWS_REGION must be set when AWS_BEARER_TOKEN_BEDROCK is configured".into(),
395                    )
396                })?;
397                let model = std::env::var("BEDROCK_MODEL")
398                    .unwrap_or_else(|_| "us.anthropic.claude-haiku-4-5-20251001-v1:0".into());
399                return Ok(Provider::Genai(GenaiProvider {
400                    kind: AdapterKind::BedrockApi,
401                    api_key: Some(api_key),
402                    endpoint: None,
403                    model,
404                    params: ModelParams::default(),
405                    headers: HashMap::new(),
406                    region: Some(region),
407                    vertex_project: None,
408                    vertex_location: None,
409                    custom_index: None,
410                }));
411            }
412        }
413
414        if let Ok(api_key) = std::env::var("ANTHROPIC_API_KEY") {
415            if !api_key.is_empty() {
416                let model = std::env::var("ANTHROPIC_MODEL")
417                    .unwrap_or_else(|_| "claude-3-5-sonnet-20241022".into());
418                return Ok(Provider::Genai(GenaiProvider {
419                    kind: AdapterKind::Anthropic,
420                    api_key: Some(api_key),
421                    endpoint: None,
422                    model,
423                    params: ModelParams::default(),
424                    headers: HashMap::new(),
425                    region: None,
426                    vertex_project: None,
427                    vertex_location: None,
428                    custom_index: None,
429                }));
430            }
431        }
432
433        if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
434            if !api_key.is_empty() {
435                let model =
436                    std::env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-2.0-flash".into());
437                return Ok(Provider::Genai(GenaiProvider {
438                    kind: AdapterKind::Gemini,
439                    api_key: Some(api_key),
440                    endpoint: None,
441                    model,
442                    params: ModelParams::default(),
443                    headers: HashMap::new(),
444                    region: None,
445                    vertex_project: None,
446                    vertex_location: None,
447                    custom_index: None,
448                }));
449            }
450        }
451
452        let base_url = std::env::var("OLLAMA_BASE_URL")
453            .or_else(|_| std::env::var("OLLAMA_URL"))
454            .unwrap_or_else(|_| "http://localhost:11434".into());
455        let model = std::env::var("OLLAMA_MODEL").unwrap_or_else(|_| "ministral-3:3b".into());
456        Ok(Provider::Genai(GenaiProvider {
457            kind: AdapterKind::Ollama,
458            api_key: None,
459            endpoint: Some(base_url),
460            model,
461            params: ModelParams::default(),
462            headers: HashMap::new(),
463            region: None,
464            vertex_project: None,
465            vertex_location: None,
466            custom_index: None,
467        }))
468    }
469
470    /// Get the provider name as a string
471    pub fn name(&self) -> &'static str {
472        match self {
473            Provider::Genai(p) => p.kind.as_lower_str(),
474            #[cfg(feature = "llamacpp")]
475            Provider::LlamaCpp { .. } => "llamacpp",
476            #[cfg(test)]
477            Provider::TestStub { .. } => "test-stub",
478        }
479    }
480
481    /// Check if this provider requires an API key
482    pub fn requires_api_key(&self) -> bool {
483        match self {
484            Provider::Genai(p) => !matches!(p.kind, AdapterKind::Ollama),
485            #[cfg(feature = "llamacpp")]
486            Provider::LlamaCpp { .. } => false,
487            #[cfg(test)]
488            Provider::TestStub { .. } => false,
489        }
490    }
491
492    /// Check if this provider is local (no network required)
493    pub fn is_local(&self) -> bool {
494        match self {
495            Provider::Genai(p) => {
496                if matches!(p.kind, AdapterKind::Ollama | AdapterKind::Omlx) {
497                    return true;
498                }
499                p.endpoint
500                    .as_deref()
501                    .map(|u| u.contains("localhost") || u.contains("127.0.0.1"))
502                    .unwrap_or(false)
503            }
504            #[cfg(feature = "llamacpp")]
505            Provider::LlamaCpp { .. } => true,
506            #[cfg(test)]
507            Provider::TestStub { .. } => true,
508        }
509    }
510
511    /// Create a provider from TOML configuration
512    pub fn from_config(
513        provider_config: &ProviderConfig,
514        model_override: Option<&str>,
515    ) -> Result<Self> {
516        Self::from_config_with_params(provider_config, model_override, ModelParams::default())
517    }
518
519    /// Create a provider from TOML configuration with model parameters
520    pub fn from_config_with_params(
521        provider_config: &ProviderConfig,
522        model_override: Option<&str>,
523        params: ModelParams,
524    ) -> Result<Self> {
525        Ok(Provider::Genai(genai_from_config(
526            provider_config,
527            model_override,
528            params,
529        )?))
530    }
531
532    /// Create a provider from a model configuration and its associated provider config
533    pub fn from_model_config(
534        model_config: &ModelConfig,
535        provider_config: &ProviderConfig,
536    ) -> Result<Self> {
537        let params = ModelParams::from_model_config(model_config);
538        Self::from_config_with_params(provider_config, Some(&model_config.model), params)
539    }
540
541    /// Create a runtime OpenAI-compatible provider from a runtime provider entry.
542    pub fn from_runtime_openai(
543        api_key: String,
544        api_base: String,
545        model: String,
546        params: ModelParams,
547        headers: HashMap<String, String>,
548    ) -> Self {
549        Provider::Genai(GenaiProvider::openai(
550            api_key, api_base, model, params, headers,
551        ))
552    }
553
554    /// Create a runtime Bedrock provider from a runtime provider entry.
555    pub fn from_runtime_bedrock(
556        api_key: String,
557        region: String,
558        model: String,
559        params: ModelParams,
560    ) -> Self {
561        Provider::Genai(GenaiProvider {
562            kind: AdapterKind::BedrockApi,
563            api_key: Some(api_key),
564            endpoint: None,
565            model,
566            params,
567            headers: HashMap::new(),
568            region: Some(region),
569            vertex_project: None,
570            vertex_location: None,
571            custom_index: None,
572        })
573    }
574}
575
576fn require_env(name: &str, what: &str) -> Result<String> {
577    std::env::var(name).map_err(|_| {
578        AppError::Configuration(format!("{what} environment variable '{name}' is not set"))
579    })
580}
581
582fn pick_model(model_override: Option<&str>, default_model: &str) -> String {
583    model_override
584        .map(String::from)
585        .unwrap_or_else(|| default_model.to_string())
586}
587
588fn genai_from_config(
589    config: &ProviderConfig,
590    model_override: Option<&str>,
591    params: ModelParams,
592) -> Result<GenaiProvider> {
593    match config {
594        ProviderConfig::OpenAI {
595            api_key_env,
596            api_base,
597            default_model,
598        } => {
599            let api_key = require_env(api_key_env, "OpenAI API key")?;
600            Ok(GenaiProvider::openai(
601                api_key,
602                api_base.clone(),
603                pick_model(model_override, default_model),
604                params,
605                HashMap::new(),
606            ))
607        }
608        ProviderConfig::Azure {
609            api_key_env,
610            base_url_env,
611            default_model,
612        } => {
613            let api_key = require_env(api_key_env, "Azure Foundry API key")?;
614            let api_base = require_env(base_url_env, "Azure Foundry base URL")?;
615            Ok(GenaiProvider::openai(
616                api_key.clone(),
617                azure_normalize_base_url(&api_base),
618                azure_strip_model_prefix(&pick_model(model_override, default_model)).to_string(),
619                params,
620                azure_foundry_headers(&api_key),
621            ))
622        }
623        ProviderConfig::Anthropic {
624            api_key_env,
625            default_model,
626        } => {
627            let api_key = require_env(api_key_env, "Anthropic API key")?;
628            Ok(GenaiProvider {
629                kind: AdapterKind::Anthropic,
630                api_key: Some(api_key),
631                endpoint: None,
632                model: pick_model(model_override, default_model),
633                params,
634                headers: HashMap::new(),
635                region: None,
636                vertex_project: None,
637                vertex_location: None,
638                custom_index: None,
639            })
640        }
641        ProviderConfig::Bedrock {
642            api_key_env,
643            region_env,
644            default_model,
645        } => {
646            let api_key = require_env(api_key_env, "Bedrock API key")?;
647            let region = require_env(region_env, "Bedrock region")?;
648            Ok(GenaiProvider {
649                kind: AdapterKind::BedrockApi,
650                api_key: Some(api_key),
651                endpoint: None,
652                model: pick_model(model_override, default_model),
653                params,
654                headers: HashMap::new(),
655                region: Some(region),
656                vertex_project: None,
657                vertex_location: None,
658                custom_index: None,
659            })
660        }
661        ProviderConfig::Ollama {
662            base_url,
663            default_model,
664            ..
665        } => Ok(GenaiProvider {
666            kind: AdapterKind::Ollama,
667            api_key: None,
668            endpoint: Some(base_url.clone()),
669            model: pick_model(model_override, default_model),
670            params,
671            headers: HashMap::new(),
672            region: None,
673            vertex_project: None,
674            vertex_location: None,
675            custom_index: None,
676        }),
677        ProviderConfig::Vertex {
678            api_key_env,
679            project_env,
680            location_env,
681            default_model,
682        } => {
683            let api_key = require_env(api_key_env, "Vertex API key")?;
684            let project = std::env::var(project_env).ok();
685            let location = std::env::var(location_env).ok();
686            Ok(GenaiProvider {
687                kind: AdapterKind::Vertex,
688                api_key: Some(api_key),
689                endpoint: None,
690                model: pick_model(model_override, default_model),
691                params,
692                headers: HashMap::new(),
693                region: None,
694                vertex_project: project,
695                vertex_location: location,
696                custom_index: None,
697            })
698        }
699        ProviderConfig::Custom {
700            index,
701            endpoint,
702            api_key_env,
703            default_model,
704        } => {
705            let api_key = match api_key_env {
706                Some(env) if !env.is_empty() => Some(require_env(env, "Custom API key")?),
707                _ => std::env::var(format!("GENAI_{index}_API_KEY")).ok(),
708            };
709            Ok(GenaiProvider {
710                kind: AdapterKind::Custom(*index),
711                api_key,
712                endpoint: Some(endpoint.clone()),
713                model: pick_model(model_override, default_model),
714                params,
715                headers: HashMap::new(),
716                region: None,
717                vertex_project: None,
718                vertex_location: None,
719                custom_index: Some(*index),
720            })
721        }
722        other => {
723            let (kind, api_key_env, api_base, default_model) = simple_genai_fields(other);
724            let optional_key = matches!(kind, AdapterKind::Omlx);
725            let api_key = if optional_key {
726                std::env::var(api_key_env).ok().filter(|s| !s.is_empty())
727            } else {
728                Some(require_env(
729                    api_key_env,
730                    &format!("{} API key", other.type_name()),
731                )?)
732            };
733            Ok(GenaiProvider {
734                kind,
735                api_key,
736                endpoint: api_base.filter(|s| !s.is_empty()),
737                model: pick_model(model_override, default_model),
738                params,
739                headers: HashMap::new(),
740                region: None,
741                vertex_project: None,
742                vertex_location: None,
743                custom_index: None,
744            })
745        }
746    }
747}
748
749fn simple_genai_fields(config: &ProviderConfig) -> (AdapterKind, &str, Option<String>, &str) {
750    match config {
751        ProviderConfig::OpenAIResp {
752            api_key_env,
753            api_base,
754            default_model,
755        } => (
756            AdapterKind::OpenAIResp,
757            api_key_env,
758            api_base.clone(),
759            default_model,
760        ),
761        ProviderConfig::Gemini {
762            api_key_env,
763            api_base,
764            default_model,
765        } => (
766            AdapterKind::Gemini,
767            api_key_env,
768            api_base.clone(),
769            default_model,
770        ),
771        ProviderConfig::Fireworks {
772            api_key_env,
773            api_base,
774            default_model,
775        } => (
776            AdapterKind::Fireworks,
777            api_key_env,
778            api_base.clone(),
779            default_model,
780        ),
781        ProviderConfig::Together {
782            api_key_env,
783            api_base,
784            default_model,
785        } => (
786            AdapterKind::Together,
787            api_key_env,
788            api_base.clone(),
789            default_model,
790        ),
791        ProviderConfig::Groq {
792            api_key_env,
793            api_base,
794            default_model,
795        } => (
796            AdapterKind::Groq,
797            api_key_env,
798            api_base.clone(),
799            default_model,
800        ),
801        ProviderConfig::Aihubmix {
802            api_key_env,
803            api_base,
804            default_model,
805        } => (
806            AdapterKind::Aihubmix,
807            api_key_env,
808            api_base.clone(),
809            default_model,
810        ),
811        ProviderConfig::Kimi {
812            api_key_env,
813            api_base,
814            default_model,
815        } => (
816            AdapterKind::Kimi,
817            api_key_env,
818            api_base.clone(),
819            default_model,
820        ),
821        ProviderConfig::Mimo {
822            api_key_env,
823            api_base,
824            default_model,
825        } => (
826            AdapterKind::Mimo,
827            api_key_env,
828            api_base.clone(),
829            default_model,
830        ),
831        ProviderConfig::Moonshot {
832            api_key_env,
833            api_base,
834            default_model,
835        } => (
836            AdapterKind::Moonshot,
837            api_key_env,
838            api_base.clone(),
839            default_model,
840        ),
841        ProviderConfig::Nebius {
842            api_key_env,
843            api_base,
844            default_model,
845        } => (
846            AdapterKind::Nebius,
847            api_key_env,
848            api_base.clone(),
849            default_model,
850        ),
851        ProviderConfig::Xai {
852            api_key_env,
853            api_base,
854            default_model,
855        } => (
856            AdapterKind::Xai,
857            api_key_env,
858            api_base.clone(),
859            default_model,
860        ),
861        ProviderConfig::DeepSeek {
862            api_key_env,
863            api_base,
864            default_model,
865        } => (
866            AdapterKind::DeepSeek,
867            api_key_env,
868            api_base.clone(),
869            default_model,
870        ),
871        ProviderConfig::Zai {
872            api_key_env,
873            api_base,
874            default_model,
875        } => (
876            AdapterKind::Zai,
877            api_key_env,
878            api_base.clone(),
879            default_model,
880        ),
881        ProviderConfig::BigModel {
882            api_key_env,
883            api_base,
884            default_model,
885        } => (
886            AdapterKind::BigModel,
887            api_key_env,
888            api_base.clone(),
889            default_model,
890        ),
891        ProviderConfig::Aliyun {
892            api_key_env,
893            api_base,
894            default_model,
895        } => (
896            AdapterKind::Aliyun,
897            api_key_env,
898            api_base.clone(),
899            default_model,
900        ),
901        ProviderConfig::QwenCloud {
902            api_key_env,
903            api_base,
904            default_model,
905        } => (
906            AdapterKind::QwenCloud,
907            api_key_env,
908            api_base.clone(),
909            default_model,
910        ),
911        ProviderConfig::Baidu {
912            api_key_env,
913            api_base,
914            default_model,
915        } => (
916            AdapterKind::Baidu,
917            api_key_env,
918            api_base.clone(),
919            default_model,
920        ),
921        ProviderConfig::Cohere {
922            api_key_env,
923            api_base,
924            default_model,
925        } => (
926            AdapterKind::Cohere,
927            api_key_env,
928            api_base.clone(),
929            default_model,
930        ),
931        ProviderConfig::OllamaCloud {
932            api_key_env,
933            api_base,
934            default_model,
935        } => (
936            AdapterKind::OllamaCloud,
937            api_key_env,
938            api_base.clone(),
939            default_model,
940        ),
941        ProviderConfig::Omlx {
942            api_key_env,
943            api_base,
944            default_model,
945        } => (
946            AdapterKind::Omlx,
947            api_key_env,
948            api_base.clone(),
949            default_model,
950        ),
951        ProviderConfig::GithubCopilot {
952            api_key_env,
953            api_base,
954            default_model,
955        } => (
956            AdapterKind::GithubCopilot,
957            api_key_env,
958            api_base.clone(),
959            default_model,
960        ),
961        ProviderConfig::OpenCodeGo {
962            api_key_env,
963            api_base,
964            default_model,
965        } => (
966            AdapterKind::OpenCodeGo,
967            api_key_env,
968            api_base.clone(),
969            default_model,
970        ),
971        ProviderConfig::BedrockApi {
972            api_key_env,
973            api_base,
974            default_model,
975        } => (
976            AdapterKind::BedrockApi,
977            api_key_env,
978            api_base.clone(),
979            default_model,
980        ),
981        ProviderConfig::OpenRouter {
982            api_key_env,
983            api_base,
984            default_model,
985        } => (
986            AdapterKind::OpenRouter,
987            api_key_env,
988            api_base.clone(),
989            default_model,
990        ),
991        ProviderConfig::AtlasCloud {
992            api_key_env,
993            api_base,
994            default_model,
995        } => (
996            AdapterKind::AtlasCloud,
997            api_key_env,
998            api_base.clone(),
999            default_model,
1000        ),
1001        ProviderConfig::MiniMax {
1002            api_key_env,
1003            api_base,
1004            default_model,
1005        } => (
1006            AdapterKind::MiniMax,
1007            api_key_env,
1008            api_base.clone(),
1009            default_model,
1010        ),
1011        other => unreachable!("simple_genai_fields on {}", other.type_name()),
1012    }
1013}
1014
1015/// Trait abstraction for LLM client factories (useful for mocking in tests)
1016#[async_trait]
1017pub trait LLMClientFactoryTrait: Send + Sync {
1018    /// Get the default provider configuration
1019    fn default_provider(&self) -> &Provider;
1020
1021    /// Create an LLM client using the default provider
1022    async fn create_default(&self) -> Result<Box<dyn LLMClient>>;
1023
1024    /// Create an LLM client using a specific provider
1025    async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>>;
1026}
1027
1028/// Configuration-based LLM client factory
1029pub struct LLMClientFactory {
1030    default_provider: Provider,
1031}
1032
1033impl LLMClientFactory {
1034    /// Create a new factory with a specific default provider
1035    pub fn new(default_provider: Provider) -> Self {
1036        Self { default_provider }
1037    }
1038
1039    /// Create a factory from environment variables
1040    pub fn from_env() -> Result<Self> {
1041        Ok(Self {
1042            default_provider: Provider::from_env()?,
1043        })
1044    }
1045
1046    /// Get the default provider configuration
1047    pub fn default_provider(&self) -> &Provider {
1048        &self.default_provider
1049    }
1050
1051    /// Create an LLM client using the default provider
1052    pub async fn create_default(&self) -> Result<Box<dyn LLMClient>> {
1053        self.default_provider.create_client().await
1054    }
1055
1056    /// Create an LLM client using a specific provider
1057    pub async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>> {
1058        provider.create_client().await
1059    }
1060}
1061
1062#[async_trait]
1063impl LLMClientFactoryTrait for LLMClientFactory {
1064    fn default_provider(&self) -> &Provider {
1065        &self.default_provider
1066    }
1067
1068    async fn create_default(&self) -> Result<Box<dyn LLMClient>> {
1069        self.default_provider.create_client().await
1070    }
1071
1072    async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>> {
1073        provider.create_client().await
1074    }
1075}
1076
1077/// Test doubles shared across crate unit tests.
1078#[cfg(test)]
1079pub(crate) mod test_support {
1080    use super::*;
1081    use ares_types::types::ToolDefinition;
1082    use async_trait::async_trait;
1083    use std::sync::atomic::{AtomicU64, Ordering};
1084
1085    /// Minimal LLM client for pool tests — never performs network I/O.
1086    pub struct MockLLMClient {
1087        model: String,
1088        id: u64,
1089    }
1090
1091    impl MockLLMClient {
1092        pub fn new(model: impl Into<String>) -> Self {
1093            static NEXT_ID: AtomicU64 = AtomicU64::new(0);
1094            Self {
1095                model: model.into(),
1096                id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
1097            }
1098        }
1099    }
1100
1101    #[async_trait]
1102    impl LLMClient for MockLLMClient {
1103        async fn generate(&self, _prompt: &str) -> Result<String> {
1104            Ok(format!("mock-{}", self.id))
1105        }
1106
1107        async fn generate_with_system(&self, _system: &str, _prompt: &str) -> Result<String> {
1108            Ok(format!("mock-{}", self.id))
1109        }
1110
1111        async fn generate_with_history(
1112            &self,
1113            _messages: &[(String, String)],
1114        ) -> Result<LLMResponse> {
1115            Ok(LLMResponse {
1116                content: format!("mock-{}", self.id),
1117                tool_calls: vec![],
1118                finish_reason: "stop".into(),
1119                usage: None,
1120                reasoning_content: None,
1121                response_id: None,
1122            })
1123        }
1124
1125        async fn generate_with_tools(
1126            &self,
1127            _prompt: &str,
1128            _tools: &[ToolDefinition],
1129        ) -> Result<LLMResponse> {
1130            Ok(LLMResponse {
1131                content: format!("mock-{}", self.id),
1132                tool_calls: vec![],
1133                finish_reason: "stop".into(),
1134                usage: None,
1135                reasoning_content: None,
1136                response_id: None,
1137            })
1138        }
1139
1140        async fn generate_with_tools_and_history(
1141            &self,
1142            _messages: &[crate::coordinator::ConversationMessage],
1143            _tools: &[ToolDefinition],
1144        ) -> Result<LLMResponse> {
1145            Ok(LLMResponse {
1146                content: format!("mock-{}", self.id),
1147                tool_calls: vec![],
1148                finish_reason: "stop".into(),
1149                usage: None,
1150                reasoning_content: None,
1151                response_id: None,
1152            })
1153        }
1154
1155        async fn stream(
1156            &self,
1157            _prompt: &str,
1158        ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1159            Err(AppError::Internal("mock stream not implemented".into()))
1160        }
1161
1162        async fn stream_with_system(
1163            &self,
1164            _system: &str,
1165            _prompt: &str,
1166        ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1167            Err(AppError::Internal("mock stream not implemented".into()))
1168        }
1169
1170        async fn stream_with_history(
1171            &self,
1172            _messages: &[(String, String)],
1173        ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1174            Err(AppError::Internal("mock stream not implemented".into()))
1175        }
1176
1177        fn model_name(&self) -> &str {
1178            &self.model
1179        }
1180    }
1181}
1182
1183#[cfg(test)]
1184mod tests {
1185    use super::*;
1186
1187    fn llm_text(content: impl Into<String>) -> LLMResponse {
1188        LLMResponse {
1189            content: content.into(),
1190            tool_calls: vec![],
1191            finish_reason: "stop".to_string(),
1192            usage: None,
1193            reasoning_content: None,
1194            response_id: None,
1195        }
1196    }
1197
1198    fn genai_openai(api_base: &str, model: &str) -> Provider {
1199        Provider::Genai(GenaiProvider::openai(
1200            "sk-test".into(),
1201            api_base.into(),
1202            model.into(),
1203            ModelParams::default(),
1204            HashMap::new(),
1205        ))
1206    }
1207
1208    #[test]
1209    fn test_llm_response_creation() {
1210        let response = llm_text("Hello");
1211        assert_eq!(response.content, "Hello");
1212        assert!(response.tool_calls.is_empty());
1213        assert_eq!(response.finish_reason, "stop");
1214        assert!(response.usage.is_none());
1215        assert!(response.reasoning_content.is_none());
1216        assert!(response.response_id.is_none());
1217    }
1218
1219    #[test]
1220    fn test_llm_response_with_usage() {
1221        let usage = TokenUsage::new(100, 50);
1222        let response = LLMResponse {
1223            content: "Hello".to_string(),
1224            tool_calls: vec![],
1225            finish_reason: "stop".to_string(),
1226            usage: Some(usage),
1227            reasoning_content: None,
1228            response_id: None,
1229        };
1230        assert!(response.usage.is_some());
1231        let usage = response.usage.unwrap();
1232        assert_eq!(usage.prompt_tokens, 100);
1233        assert_eq!(usage.completion_tokens, 50);
1234        assert_eq!(usage.total_tokens, 150);
1235    }
1236
1237    #[test]
1238    fn test_llm_response_with_tool_calls() {
1239        let tool_calls = vec![
1240            ToolCall {
1241                id: "1".to_string(),
1242                name: "calculator".to_string(),
1243                arguments: serde_json::json!({"a": 1, "b": 2}),
1244            },
1245            ToolCall {
1246                id: "2".to_string(),
1247                name: "search".to_string(),
1248                arguments: serde_json::json!({"query": "test"}),
1249            },
1250        ];
1251
1252        let response = LLMResponse {
1253            content: "".to_string(),
1254            tool_calls,
1255            finish_reason: "tool_calls".to_string(),
1256            usage: Some(TokenUsage::new(50, 25)),
1257            reasoning_content: None,
1258            response_id: None,
1259        };
1260
1261        assert_eq!(response.tool_calls.len(), 2);
1262        assert_eq!(response.tool_calls[0].name, "calculator");
1263        assert_eq!(response.finish_reason, "tool_calls");
1264        assert_eq!(response.usage.as_ref().unwrap().total_tokens, 75);
1265    }
1266
1267    #[test]
1268    fn test_factory_creation() {
1269        let factory = LLMClientFactory::new(genai_openai("https://api.openai.com/v1", "test"));
1270        assert_eq!(factory.default_provider().name(), "openai");
1271    }
1272
1273    #[test]
1274    fn test_openai_provider_properties() {
1275        let provider = genai_openai("https://api.openai.com/v1", "gpt-4");
1276        assert_eq!(provider.name(), "openai");
1277        assert!(provider.requires_api_key());
1278        assert!(!provider.is_local());
1279    }
1280
1281    #[test]
1282    fn test_openai_local_provider() {
1283        let provider = genai_openai("http://localhost:8000/v1", "local-model");
1284        assert!(provider.is_local());
1285    }
1286
1287    #[test]
1288    fn test_token_usage_default_all_zeros() {
1289        let usage = TokenUsage::default();
1290        assert_eq!(usage.prompt_tokens, 0);
1291        assert_eq!(usage.completion_tokens, 0);
1292        assert_eq!(usage.total_tokens, 0);
1293    }
1294
1295    #[test]
1296    fn test_token_usage_new_calculates_total() {
1297        let usage = TokenUsage::new(100, 50);
1298        assert_eq!(usage.prompt_tokens, 100);
1299        assert_eq!(usage.completion_tokens, 50);
1300        assert_eq!(usage.total_tokens, 150);
1301    }
1302
1303    #[test]
1304    fn test_token_usage_new_zero_tokens() {
1305        let usage = TokenUsage::new(0, 0);
1306        assert_eq!(usage.total_tokens, 0);
1307    }
1308
1309    #[test]
1310    fn test_token_usage_new_large_values() {
1311        let usage = TokenUsage::new(u32::MAX / 2, u32::MAX / 2 + 1);
1312        assert_eq!(usage.total_tokens, u32::MAX);
1313    }
1314
1315    #[test]
1316    fn test_token_usage_serde_roundtrip() {
1317        let usage = TokenUsage::new(100, 200);
1318        let json = serde_json::to_string(&usage).unwrap();
1319        let deserialized: TokenUsage = serde_json::from_str(&json).unwrap();
1320        assert_eq!(usage, deserialized);
1321    }
1322
1323    #[test]
1324    fn test_token_usage_serde_default_values() {
1325        let json = r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}"#;
1326        let usage: TokenUsage = serde_json::from_str(json).unwrap();
1327        assert_eq!(usage, TokenUsage::default());
1328    }
1329
1330    #[test]
1331    fn test_token_usage_serde_partial_json() {
1332        let json = r#"{"prompt_tokens":42,"completion_tokens":58,"total_tokens":100}"#;
1333        let usage: TokenUsage = serde_json::from_str(json).unwrap();
1334        assert_eq!(usage.prompt_tokens, 42);
1335        assert_eq!(usage.completion_tokens, 58);
1336        assert_eq!(usage.total_tokens, 100);
1337    }
1338
1339    #[test]
1340    fn test_token_usage_clone_eq() {
1341        let a = TokenUsage::new(10, 20);
1342        let b = a.clone();
1343        assert_eq!(a, b);
1344    }
1345
1346    #[test]
1347    fn test_token_usage_debug_format() {
1348        let usage = TokenUsage::new(1, 2);
1349        let debug_str = format!("{:?}", usage);
1350        assert!(debug_str.contains("TokenUsage"));
1351        assert!(debug_str.contains("prompt_tokens"));
1352    }
1353
1354    #[test]
1355    fn test_model_params_default_all_none() {
1356        let params = ModelParams::default();
1357        assert!(params.temperature.is_none());
1358        assert!(params.max_tokens.is_none());
1359        assert!(params.top_p.is_none());
1360        assert!(params.frequency_penalty.is_none());
1361        assert!(params.presence_penalty.is_none());
1362    }
1363
1364    #[test]
1365    fn test_model_params_from_model_config_all_fields() {
1366        let config = ModelConfig {
1367            provider: "openai".to_string(),
1368            model: "gpt-4".to_string(),
1369            temperature: 0.5,
1370            max_tokens: 1024,
1371        };
1372        let params = ModelParams::from_model_config(&config);
1373        assert_eq!(params.temperature, Some(0.5));
1374        assert_eq!(params.max_tokens, Some(1024));
1375        assert!(params.top_p.is_none());
1376        assert!(params.frequency_penalty.is_none());
1377        assert!(params.presence_penalty.is_none());
1378    }
1379
1380    #[test]
1381    fn test_model_params_from_model_config_optional_none() {
1382        let config = ModelConfig {
1383            provider: "openai".to_string(),
1384            model: "mistral".to_string(),
1385            temperature: 0.7,
1386            max_tokens: 512,
1387        };
1388        let params = ModelParams::from_model_config(&config);
1389        assert_eq!(params.temperature, Some(0.7));
1390        assert_eq!(params.max_tokens, Some(512));
1391        assert!(params.top_p.is_none());
1392        assert!(params.frequency_penalty.is_none());
1393        assert!(params.presence_penalty.is_none());
1394    }
1395
1396    #[test]
1397    fn test_model_params_clone() {
1398        let params = ModelParams {
1399            temperature: Some(0.8),
1400            max_tokens: Some(2048),
1401            top_p: Some(0.95),
1402            frequency_penalty: Some(-0.5),
1403            presence_penalty: Some(0.3),
1404        };
1405        let cloned = params.clone();
1406        assert_eq!(params.temperature, cloned.temperature);
1407        assert_eq!(params.max_tokens, cloned.max_tokens);
1408        assert_eq!(params.top_p, cloned.top_p);
1409        assert_eq!(params.frequency_penalty, cloned.frequency_penalty);
1410        assert_eq!(params.presence_penalty, cloned.presence_penalty);
1411    }
1412
1413    #[test]
1414    fn test_llm_response_empty_content() {
1415        let response = llm_text("");
1416        assert!(response.content.is_empty());
1417    }
1418
1419    #[test]
1420    fn test_llm_response_clone() {
1421        let response = LLMResponse {
1422            content: "hello".to_string(),
1423            tool_calls: vec![ToolCall {
1424                id: "1".to_string(),
1425                name: "fn".to_string(),
1426                arguments: serde_json::json!({"key": "value"}),
1427            }],
1428            finish_reason: "tool_calls".to_string(),
1429            usage: Some(TokenUsage::new(10, 20)),
1430            reasoning_content: None,
1431            response_id: None,
1432        };
1433        let cloned = response.clone();
1434        assert_eq!(cloned.content, "hello");
1435        assert_eq!(cloned.tool_calls.len(), 1);
1436        assert_eq!(cloned.tool_calls[0].name, "fn");
1437        assert_eq!(cloned.finish_reason, "tool_calls");
1438        assert_eq!(cloned.usage.unwrap().total_tokens, 30);
1439    }
1440
1441    #[test]
1442    fn test_openai_from_config_missing_env_var() {
1443        std::env::remove_var("TEST_OPENAI_MISSING_KEY");
1444        let config = ProviderConfig::OpenAI {
1445            api_key_env: "TEST_OPENAI_MISSING_KEY".to_string(),
1446            api_base: "https://api.openai.com/v1".to_string(),
1447            default_model: "gpt-4".to_string(),
1448        };
1449        let result = Provider::from_config(&config, None);
1450        assert!(result.is_err());
1451        match result.unwrap_err() {
1452            AppError::Configuration(msg) => {
1453                assert!(msg.contains("TEST_OPENAI_MISSING_KEY"));
1454            }
1455            other => panic!("Expected Configuration error, got: {:?}", other),
1456        }
1457    }
1458
1459    #[test]
1460    fn test_token_usage_not_equal() {
1461        assert_ne!(TokenUsage::new(1, 2), TokenUsage::new(3, 4));
1462    }
1463
1464    #[test]
1465    fn test_model_params_debug_format() {
1466        let params = ModelParams::default();
1467        let debug_str = format!("{:?}", params);
1468        assert!(debug_str.contains("ModelParams"));
1469    }
1470
1471    fn test_stub_provider(model: &str) -> Provider {
1472        Provider::TestStub {
1473            model: model.to_string(),
1474        }
1475    }
1476
1477    #[test]
1478    fn test_stub_provider_properties() {
1479        let provider = test_stub_provider("unit-test");
1480        assert_eq!(provider.name(), "test-stub");
1481        assert!(!provider.requires_api_key());
1482        assert!(provider.is_local());
1483    }
1484
1485    #[tokio::test]
1486    async fn test_provider_create_client_test_stub() {
1487        let client = test_stub_provider("provider-model")
1488            .create_client()
1489            .await
1490            .expect("TestStub client");
1491        assert_eq!(client.model_name(), "provider-model");
1492    }
1493
1494    #[tokio::test]
1495    async fn test_factory_create_default_via_test_stub() {
1496        let factory = LLMClientFactory::new(test_stub_provider("factory-model"));
1497        let client = factory.create_default().await.expect("factory client");
1498        assert_eq!(client.model_name(), "factory-model");
1499    }
1500
1501    #[tokio::test]
1502    async fn test_factory_trait_create_with_provider() {
1503        let factory = LLMClientFactory::new(test_stub_provider("default"));
1504        let trait_ref: &dyn LLMClientFactoryTrait = &factory;
1505        let client = trait_ref
1506            .create_with_provider(test_stub_provider("switched"))
1507            .await
1508            .expect("switched client");
1509        assert_eq!(client.model_name(), "switched");
1510    }
1511
1512    mod llm_client_trait_tests {
1513        use super::*;
1514        use crate::client::test_support::MockLLMClient;
1515        use crate::coordinator::ConversationMessage;
1516        use ares_types::types::ToolDefinition;
1517
1518        #[tokio::test]
1519        async fn test_generate_and_model_name() {
1520            let client = MockLLMClient::new("trait-model");
1521            assert_eq!(client.model_name(), "trait-model");
1522            let out = client.generate("hello").await.expect("generate");
1523            assert!(out.starts_with("mock-"));
1524        }
1525
1526        #[tokio::test]
1527        async fn test_generate_with_system() {
1528            let client = MockLLMClient::new("sys");
1529            let out = client
1530                .generate_with_system("system", "prompt")
1531                .await
1532                .expect("generate_with_system");
1533            assert!(out.starts_with("mock-"));
1534        }
1535
1536        #[tokio::test]
1537        async fn test_generate_with_history() {
1538            let client = MockLLMClient::new("hist");
1539            let messages = vec![("user".to_string(), "hi".to_string())];
1540            let response = client
1541                .generate_with_history(&messages)
1542                .await
1543                .expect("generate_with_history");
1544            assert!(response.content.starts_with("mock-"));
1545            assert_eq!(response.finish_reason, "stop");
1546            assert!(response.tool_calls.is_empty());
1547        }
1548
1549        #[tokio::test]
1550        async fn test_generate_with_tools() {
1551            let client = MockLLMClient::new("tools");
1552            let tools = vec![ToolDefinition {
1553                name: "search".to_string(),
1554                description: "Search".to_string(),
1555                parameters: serde_json::json!({"type": "object"}),
1556            }];
1557            let response = client
1558                .generate_with_tools("find docs", &tools)
1559                .await
1560                .expect("generate_with_tools");
1561            assert!(response.content.starts_with("mock-"));
1562        }
1563
1564        #[tokio::test]
1565        async fn test_generate_with_tools_and_history() {
1566            let client = MockLLMClient::new("both");
1567            let messages = vec![ConversationMessage::user("run tool")];
1568            let tools = vec![ToolDefinition {
1569                name: "calc".to_string(),
1570                description: "Calculate".to_string(),
1571                parameters: serde_json::json!({"type": "object"}),
1572            }];
1573            let response = client
1574                .generate_with_tools_and_history(&messages, &tools)
1575                .await
1576                .expect("generate_with_tools_and_history");
1577            assert!(response.content.starts_with("mock-"));
1578        }
1579
1580        #[tokio::test]
1581        async fn test_stream_methods_return_internal_error() {
1582            let client = MockLLMClient::new("stream");
1583            for result in [
1584                client.stream("hi").await,
1585                client.stream_with_system("sys", "hi").await,
1586                client
1587                    .stream_with_history(&[("user".into(), "hi".into())])
1588                    .await,
1589            ] {
1590                assert!(matches!(result, Err(AppError::Internal(_))));
1591            }
1592        }
1593
1594        #[test]
1595        fn default_supports_hints_is_false() {
1596            let client = MockLLMClient::new("hints");
1597            assert!(!client.supports_hints());
1598            client.set_hints(GenerationHints {
1599                json_mode: true,
1600                ..Default::default()
1601            });
1602        }
1603
1604        #[test]
1605        fn hint_recording_mock_records_set_hints_calls() {
1606            use parking_lot::Mutex;
1607            use std::sync::Arc;
1608
1609            #[derive(Default)]
1610            struct HintRecordingClient {
1611                hints: Mutex<Vec<GenerationHints>>,
1612            }
1613
1614            #[async_trait]
1615            impl LLMClient for HintRecordingClient {
1616                async fn generate(&self, _prompt: &str) -> Result<String> {
1617                    Err(AppError::Internal("unused".into()))
1618                }
1619
1620                async fn generate_with_system(
1621                    &self,
1622                    _system: &str,
1623                    _prompt: &str,
1624                ) -> Result<String> {
1625                    Err(AppError::Internal("unused".into()))
1626                }
1627
1628                async fn generate_with_history(
1629                    &self,
1630                    _messages: &[(String, String)],
1631                ) -> Result<LLMResponse> {
1632                    Err(AppError::Internal("unused".into()))
1633                }
1634
1635                async fn generate_with_tools(
1636                    &self,
1637                    _prompt: &str,
1638                    _tools: &[ToolDefinition],
1639                ) -> Result<LLMResponse> {
1640                    Err(AppError::Internal("unused".into()))
1641                }
1642
1643                async fn generate_with_tools_and_history(
1644                    &self,
1645                    _messages: &[crate::coordinator::ConversationMessage],
1646                    _tools: &[ToolDefinition],
1647                ) -> Result<LLMResponse> {
1648                    Err(AppError::Internal("unused".into()))
1649                }
1650
1651                async fn stream(
1652                    &self,
1653                    _prompt: &str,
1654                ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1655                {
1656                    Err(AppError::Internal("unused".into()))
1657                }
1658
1659                async fn stream_with_system(
1660                    &self,
1661                    _system: &str,
1662                    _prompt: &str,
1663                ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1664                {
1665                    Err(AppError::Internal("unused".into()))
1666                }
1667
1668                async fn stream_with_history(
1669                    &self,
1670                    _messages: &[(String, String)],
1671                ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1672                {
1673                    Err(AppError::Internal("unused".into()))
1674                }
1675
1676                fn model_name(&self) -> &str {
1677                    "hint-recording-mock"
1678                }
1679
1680                fn supports_hints(&self) -> bool {
1681                    true
1682                }
1683
1684                fn set_hints(&self, hints: GenerationHints) {
1685                    self.hints.lock().push(hints);
1686                }
1687            }
1688
1689            let client = Arc::new(HintRecordingClient::default());
1690            assert!(client.supports_hints());
1691            client.set_hints(GenerationHints {
1692                json_mode: true,
1693                suppress_reasoning: false,
1694                max_tokens: Some(256),
1695                guided_grammar: None,
1696                ..Default::default()
1697            });
1698            client.set_hints(GenerationHints::default());
1699
1700            let recorded = client.hints.lock();
1701            assert_eq!(
1702                recorded.len(),
1703                2,
1704                "every set_hints call is recorded in order"
1705            );
1706            assert!(recorded[0].json_mode && recorded[0].max_tokens == Some(256));
1707            assert_eq!(
1708                recorded[1],
1709                GenerationHints::default(),
1710                "clearing via Default::default() reaches the impl"
1711            );
1712        }
1713    }
1714}
ares_llm/client.rs

ares_llm/
client.rs
