1use crate::config::{ModelConfig, ProviderConfig};
2use ares_types::types::{AppError, Result, ToolCall, ToolDefinition};
3use async_trait::async_trait;
4
5#[derive(Debug, Clone, Default, PartialEq, Eq)]
39pub struct GenerationHints {
40 pub json_mode: bool,
42 pub suppress_reasoning: bool,
44 pub max_tokens: Option<u32>,
46 pub guided_grammar: Option<String>,
51}
52
53#[async_trait]
55pub trait LLMClient: Send + Sync {
56 async fn generate(&self, prompt: &str) -> Result<String>;
58
59 async fn generate_with_system(&self, system: &str, prompt: &str) -> Result<String>;
61
62 async fn generate_with_history(
64 &self,
65 messages: &[(String, String)], ) -> Result<LLMResponse>;
67
68 async fn generate_with_tools(
70 &self,
71 prompt: &str,
72 tools: &[ToolDefinition],
73 ) -> Result<LLMResponse>;
74
75 async fn generate_with_tools_and_history(
90 &self,
91 messages: &[crate::coordinator::ConversationMessage],
92 tools: &[ToolDefinition],
93 ) -> Result<LLMResponse>;
94
95 async fn stream(
97 &self,
98 prompt: &str,
99 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
100
101 async fn stream_with_system(
103 &self,
104 system: &str,
105 prompt: &str,
106 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
107
108 async fn stream_with_history(
110 &self,
111 messages: &[(String, String)], ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>;
113
114 fn model_name(&self) -> &str;
116
117 fn supports_hints(&self) -> bool {
121 false
122 }
123
124 fn set_hints(&self, _hints: GenerationHints) {}
130}
131
132#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
134pub struct TokenUsage {
135 pub prompt_tokens: u32,
137 pub completion_tokens: u32,
139 pub total_tokens: u32,
141 #[serde(default)]
145 pub cached_tokens: Option<i64>,
146}
147
148impl TokenUsage {
149 pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
151 Self {
152 prompt_tokens,
153 completion_tokens,
154 total_tokens: prompt_tokens + completion_tokens,
155 cached_tokens: None,
156 }
157 }
158}
159
160#[derive(Debug, Clone)]
162pub struct LLMResponse {
163 pub content: String,
165 pub tool_calls: Vec<ToolCall>,
167 pub finish_reason: String,
169 pub usage: Option<TokenUsage>,
171}
172
173#[derive(Debug, Clone, PartialEq, Default)]
175pub struct ModelParams {
176 pub temperature: Option<f32>,
178 pub max_tokens: Option<u32>,
180 pub top_p: Option<f32>,
182 pub frequency_penalty: Option<f32>,
184 pub presence_penalty: Option<f32>,
186}
187
188impl ModelParams {
189 pub fn from_model_config(config: &ModelConfig) -> Self {
191 Self {
192 temperature: Some(config.temperature),
193 max_tokens: Some(config.max_tokens),
194 top_p: None,
195 frequency_penalty: None,
196 presence_penalty: None,
197 }
198 }
199}
200
201#[derive(Debug, Clone)]
206#[non_exhaustive]
207pub enum Provider {
208 #[cfg(feature = "openai")]
210 OpenAI {
211 api_key: String,
213 api_base: String,
215 model: String,
217 params: ModelParams,
219 },
220
221 #[cfg(feature = "azure")]
223 Azure {
224 api_key: String,
226 api_base: String,
228 model: String,
230 params: ModelParams,
232 },
233
234 #[cfg(feature = "anthropic")]
236 Anthropic {
237 api_key: String,
239 model: String,
241 params: ModelParams,
243 },
244
245 #[cfg(feature = "bedrock")]
247 Bedrock {
248 api_key: String,
250 region: String,
252 model: String,
254 params: ModelParams,
256 },
257
258 #[cfg(feature = "openai")]
260 RuntimeOpenAI {
261 api_key: String,
263 api_base: String,
265 model: String,
267 params: ModelParams,
269 headers: std::collections::HashMap<String, String>,
271 },
272
273 #[cfg(feature = "ollama")]
275 Ollama {
276 base_url: String,
278 model: String,
280 params: ModelParams,
282 },
283
284 #[cfg(test)]
286 TestStub {
287 model: String,
289 },
290}
291
292impl Provider {
293 pub async fn create_client(&self) -> Result<Box<dyn LLMClient>> {
302 match self {
303 #[cfg(feature = "openai")]
304 Provider::OpenAI {
305 api_key,
306 api_base,
307 model,
308 params,
309 } => Ok(Box::new(super::openai::OpenAIClient::with_params(
310 api_key.clone(),
311 api_base.clone(),
312 model.clone(),
313 params.clone(),
314 ))),
315
316 #[cfg(feature = "azure")]
317 Provider::Azure {
318 api_key,
319 api_base,
320 model,
321 params,
322 } => Ok(Box::new(
323 super::openai::OpenAIClient::with_params_and_headers(
324 api_key.clone(),
325 super::azure::normalize_base_url(api_base),
326 super::azure::strip_model_prefix(model).to_string(),
327 params.clone(),
328 super::azure::foundry_headers(api_key),
329 ),
330 )),
331
332 #[cfg(feature = "openai")]
333 Provider::RuntimeOpenAI {
334 api_key,
335 api_base,
336 model,
337 params,
338 headers,
339 } => Ok(Box::new(
340 super::openai::OpenAIClient::with_params_and_headers(
341 api_key.clone(),
342 api_base.clone(),
343 model.clone(),
344 params.clone(),
345 headers.clone(),
346 ),
347 )),
348
349 #[cfg(feature = "anthropic")]
350 Provider::Anthropic {
351 api_key,
352 model,
353 params,
354 } => Ok(Box::new(super::anthropic::AnthropicClient::with_params(
355 api_key.clone(),
356 model.clone(),
357 params.clone(),
358 ))),
359
360 #[cfg(feature = "bedrock")]
361 Provider::Bedrock {
362 api_key,
363 region,
364 model,
365 params,
366 } => Ok(Box::new(super::bedrock::BedrockClient::with_params(
367 api_key.clone(),
368 region.clone(),
369 model.clone(),
370 params.clone(),
371 ))),
372
373 #[cfg(feature = "ollama")]
374 Provider::Ollama {
375 base_url,
376 model,
377 params,
378 } => super::ollama::OllamaClient::with_params(
379 base_url.clone(),
380 model.clone(),
381 params.clone(),
382 )
383 .await
384 .map(|c| Box::new(c) as Box<dyn LLMClient>),
385
386 #[cfg(test)]
387 Provider::TestStub { model } => {
388 Ok(Box::new(test_support::MockLLMClient::new(model.clone())))
389 }
390
391 #[allow(unreachable_patterns)]
392 _ => Err(AppError::Configuration(
393 "No matching LLM provider feature is enabled for this provider".into(),
394 )),
395 }
396 }
397
398 pub fn from_env() -> Result<Self> {
446 #[cfg(feature = "openai")]
447 {
448 if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
449 if !api_key.is_empty() {
450 let api_base = std::env::var("OPENAI_API_BASE")
451 .unwrap_or_else(|_| "https://api.openai.com/v1".into());
452 let model = std::env::var("OPENAI_MODEL").unwrap_or_else(|_| "gpt-4".into());
453 return Ok(Provider::OpenAI {
454 api_key,
455 api_base,
456 model,
457 params: ModelParams::default(),
458 });
459 }
460 }
461
462 if let Ok(api_key) = std::env::var("NVIDIA_API_KEY") {
464 if !api_key.is_empty() {
465 return Ok(Provider::OpenAI {
466 api_key,
467 api_base: "https://integrate.api.nvidia.com/v1".into(),
468 model: "nvidia/nemotron-3-ultra-550b-a55b".into(),
469 params: ModelParams::default(),
470 });
471 }
472 }
473 }
474
475 #[cfg(feature = "azure")]
476 {
477 if let Ok(api_key) = std::env::var(super::azure::DEFAULT_API_KEY_ENV) {
478 if !api_key.is_empty() {
479 let api_base =
480 std::env::var(super::azure::DEFAULT_BASE_URL_ENV).map_err(|_| {
481 AppError::Configuration(format!(
482 "{} must be set when {} is configured",
483 super::azure::DEFAULT_BASE_URL_ENV,
484 super::azure::DEFAULT_API_KEY_ENV
485 ))
486 })?;
487 let model = std::env::var(super::azure::DEFAULT_MODEL_ENV)
488 .unwrap_or_else(|_| super::azure::DEFAULT_MODEL.to_string());
489 return Ok(Provider::Azure {
490 api_key,
491 api_base,
492 model,
493 params: ModelParams::default(),
494 });
495 }
496 }
497 }
498
499 #[cfg(feature = "bedrock")]
500 {
501 if let Ok(api_key) = std::env::var("AWS_BEARER_TOKEN_BEDROCK") {
502 if !api_key.is_empty() {
503 let region = std::env::var("AWS_REGION").map_err(|_| {
504 AppError::Configuration(
505 "AWS_REGION must be set when AWS_BEARER_TOKEN_BEDROCK is configured"
506 .into(),
507 )
508 })?;
509 let model = std::env::var("BEDROCK_MODEL")
510 .unwrap_or_else(|_| "us.anthropic.claude-haiku-4-5-20251001-v1:0".into());
511 return Ok(Provider::Bedrock {
512 api_key,
513 region,
514 model,
515 params: ModelParams::default(),
516 });
517 }
518 }
519 }
520
521 #[cfg(all(
522 not(feature = "openai"),
523 not(feature = "azure"),
524 not(feature = "bedrock")
525 ))]
526 return Err(AppError::Configuration(
527 "No LLM provider feature is enabled. Enable openai, azure, or bedrock.".into(),
528 ));
529
530 #[cfg(any(feature = "openai", feature = "azure", feature = "bedrock"))]
531 Err(AppError::Configuration(
532 "No LLM provider configured. Set OPENAI_API_KEY, NVIDIA_API_KEY, AZURE_FOUNDRY_API_KEY, or AWS_BEARER_TOKEN_BEDROCK.".into(),
533 ))
534 }
535
536 pub fn name(&self) -> &'static str {
538 match self {
539 #[cfg(feature = "openai")]
540 Provider::OpenAI { .. } => "openai",
541
542 #[cfg(feature = "azure")]
543 Provider::Azure { .. } => "azure",
544
545 #[cfg(feature = "openai")]
546 Provider::RuntimeOpenAI { .. } => "openai",
547
548 #[cfg(feature = "anthropic")]
549 Provider::Anthropic { .. } => "anthropic",
550
551 #[cfg(feature = "bedrock")]
552 Provider::Bedrock { .. } => "bedrock",
553
554 #[cfg(feature = "ollama")]
555 Provider::Ollama { .. } => "ollama",
556
557 #[cfg(test)]
558 Provider::TestStub { .. } => "test-stub",
559
560 #[allow(unreachable_patterns)]
561 _ => "unknown",
562 }
563 }
564
565 pub fn requires_api_key(&self) -> bool {
567 match self {
568 #[cfg(feature = "openai")]
569 Provider::OpenAI { .. } => true,
570
571 #[cfg(feature = "azure")]
572 Provider::Azure { .. } => true,
573
574 #[cfg(feature = "openai")]
575 Provider::RuntimeOpenAI { .. } => true,
576
577 #[cfg(feature = "anthropic")]
578 Provider::Anthropic { .. } => true,
579
580 #[cfg(feature = "bedrock")]
581 Provider::Bedrock { .. } => true,
582
583 #[cfg(feature = "ollama")]
584 Provider::Ollama { .. } => false,
585
586 #[cfg(test)]
587 Provider::TestStub { .. } => false,
588
589 #[allow(unreachable_patterns)]
590 _ => false,
591 }
592 }
593
594 pub fn is_local(&self) -> bool {
596 match self {
597 #[cfg(feature = "openai")]
598 Provider::OpenAI { api_base, .. } => {
599 api_base.contains("localhost") || api_base.contains("127.0.0.1")
600 }
601
602 #[cfg(feature = "azure")]
603 Provider::Azure { .. } => false,
604
605 #[cfg(feature = "openai")]
606 Provider::RuntimeOpenAI { api_base, .. } => {
607 api_base.contains("localhost") || api_base.contains("127.0.0.1")
608 }
609
610 #[cfg(feature = "ollama")]
611 Provider::Ollama { base_url, .. } => {
612 base_url.contains("localhost") || base_url.contains("127.0.0.1")
613 }
614
615 #[cfg(feature = "anthropic")]
616 Provider::Anthropic { .. } => false,
617
618 #[cfg(feature = "bedrock")]
619 Provider::Bedrock { .. } => false,
620
621 #[cfg(test)]
622 Provider::TestStub { .. } => true,
623
624 #[allow(unreachable_patterns)]
625 _ => false,
626 }
627 }
628
629 #[allow(unused_variables)]
641 pub fn from_config(
642 provider_config: &ProviderConfig,
643 model_override: Option<&str>,
644 ) -> Result<Self> {
645 Self::from_config_with_params(provider_config, model_override, ModelParams::default())
646 }
647
648 #[allow(unused_variables)]
650 pub fn from_config_with_params(
651 provider_config: &ProviderConfig,
652 model_override: Option<&str>,
653 params: ModelParams,
654 ) -> Result<Self> {
655 match provider_config {
656 #[cfg(feature = "openai")]
657 ProviderConfig::OpenAI {
658 api_key_env,
659 api_base,
660 default_model,
661 } => {
662 let api_key = std::env::var(api_key_env).map_err(|_| {
663 AppError::Configuration(format!(
664 "OpenAI API key environment variable '{}' is not set",
665 api_key_env
666 ))
667 })?;
668 Ok(Provider::OpenAI {
669 api_key,
670 api_base: api_base.clone(),
671 model: model_override
672 .map(String::from)
673 .unwrap_or_else(|| default_model.clone()),
674 params,
675 })
676 }
677
678 #[cfg(feature = "azure")]
679 ProviderConfig::Azure {
680 api_key_env,
681 base_url_env,
682 default_model,
683 } => {
684 let api_key = std::env::var(api_key_env).map_err(|_| {
685 AppError::Configuration(format!(
686 "Azure Foundry API key environment variable '{}' is not set",
687 api_key_env
688 ))
689 })?;
690 let api_base = std::env::var(base_url_env).map_err(|_| {
691 AppError::Configuration(format!(
692 "Azure Foundry base URL environment variable '{}' is not set",
693 base_url_env
694 ))
695 })?;
696 Ok(Provider::Azure {
697 api_key,
698 api_base,
699 model: model_override
700 .map(String::from)
701 .unwrap_or_else(|| default_model.clone()),
702 params,
703 })
704 }
705
706 #[cfg(feature = "anthropic")]
707 ProviderConfig::Anthropic {
708 api_key_env,
709 default_model,
710 } => {
711 let api_key = std::env::var(api_key_env).map_err(|_| {
712 AppError::Configuration(format!(
713 "Anthropic API key environment variable '{}' is not set",
714 api_key_env
715 ))
716 })?;
717 Ok(Provider::Anthropic {
718 api_key,
719 model: model_override
720 .map(String::from)
721 .unwrap_or_else(|| default_model.clone()),
722 params,
723 })
724 }
725
726 #[cfg(feature = "bedrock")]
727 ProviderConfig::Bedrock {
728 api_key_env,
729 region_env,
730 default_model,
731 } => {
732 let api_key = std::env::var(api_key_env).map_err(|_| {
733 AppError::Configuration(format!(
734 "Bedrock API key environment variable '{}' is not set",
735 api_key_env
736 ))
737 })?;
738 let region = std::env::var(region_env).map_err(|_| {
739 AppError::Configuration(format!(
740 "Bedrock region environment variable '{}' is not set",
741 region_env
742 ))
743 })?;
744 Ok(Provider::Bedrock {
745 api_key,
746 region,
747 model: model_override
748 .map(String::from)
749 .unwrap_or_else(|| default_model.clone()),
750 params,
751 })
752 }
753
754 #[cfg(feature = "ollama")]
755 ProviderConfig::Ollama {
756 base_url,
757 default_model,
758 ..
759 } => Ok(Provider::Ollama {
760 base_url: base_url.clone(),
761 model: model_override
762 .map(String::from)
763 .unwrap_or_else(|| default_model.clone()),
764 params,
765 }),
766
767 #[allow(unreachable_patterns)]
770 _ => Err(AppError::Configuration(format!(
771 "{} provider configured but the corresponding feature is not enabled in this build",
772 provider_config.type_name()
773 ))),
774 }
775 }
776
777 pub fn from_model_config(
782 model_config: &ModelConfig,
783 provider_config: &ProviderConfig,
784 ) -> Result<Self> {
785 let params = ModelParams::from_model_config(model_config);
786 Self::from_config_with_params(provider_config, Some(&model_config.model), params)
787 }
788
789 #[cfg(feature = "openai")]
791 pub fn from_runtime_openai(
792 api_key: String,
793 api_base: String,
794 model: String,
795 params: ModelParams,
796 headers: std::collections::HashMap<String, String>,
797 ) -> Self {
798 Provider::RuntimeOpenAI {
799 api_key,
800 api_base,
801 model,
802 params,
803 headers,
804 }
805 }
806
807 #[cfg(feature = "bedrock")]
809 pub fn from_runtime_bedrock(
810 api_key: String,
811 region: String,
812 model: String,
813 params: ModelParams,
814 ) -> Self {
815 Provider::Bedrock {
816 api_key,
817 region,
818 model,
819 params,
820 }
821 }
822}
823
824#[async_trait]
826pub trait LLMClientFactoryTrait: Send + Sync {
827 fn default_provider(&self) -> &Provider;
829
830 async fn create_default(&self) -> Result<Box<dyn LLMClient>>;
832
833 async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>>;
835}
836
837pub struct LLMClientFactory {
842 default_provider: Provider,
843}
844
845impl LLMClientFactory {
846 pub fn new(default_provider: Provider) -> Self {
848 Self { default_provider }
849 }
850
851 pub fn from_env() -> Result<Self> {
855 Ok(Self {
856 default_provider: Provider::from_env()?,
857 })
858 }
859
860 pub fn default_provider(&self) -> &Provider {
862 &self.default_provider
863 }
864
865 pub async fn create_default(&self) -> Result<Box<dyn LLMClient>> {
867 self.default_provider.create_client().await
868 }
869
870 pub async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>> {
872 provider.create_client().await
873 }
874}
875
876#[async_trait]
877impl LLMClientFactoryTrait for LLMClientFactory {
878 fn default_provider(&self) -> &Provider {
879 &self.default_provider
880 }
881
882 async fn create_default(&self) -> Result<Box<dyn LLMClient>> {
883 self.default_provider.create_client().await
884 }
885
886 async fn create_with_provider(&self, provider: Provider) -> Result<Box<dyn LLMClient>> {
887 provider.create_client().await
888 }
889}
890
891#[cfg(test)]
893pub(crate) mod test_support {
894 use super::*;
895 use ares_types::types::ToolDefinition;
896 use async_trait::async_trait;
897 use std::sync::atomic::{AtomicU64, Ordering};
898
899 pub struct MockLLMClient {
901 model: String,
902 id: u64,
903 }
904
905 impl MockLLMClient {
906 pub fn new(model: impl Into<String>) -> Self {
907 static NEXT_ID: AtomicU64 = AtomicU64::new(0);
908 Self {
909 model: model.into(),
910 id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
911 }
912 }
913 }
914
915 #[async_trait]
916 impl LLMClient for MockLLMClient {
917 async fn generate(&self, _prompt: &str) -> Result<String> {
918 Ok(format!("mock-{}", self.id))
919 }
920
921 async fn generate_with_system(&self, _system: &str, _prompt: &str) -> Result<String> {
922 Ok(format!("mock-{}", self.id))
923 }
924
925 async fn generate_with_history(
926 &self,
927 _messages: &[(String, String)],
928 ) -> Result<LLMResponse> {
929 Ok(LLMResponse {
930 content: format!("mock-{}", self.id),
931 tool_calls: vec![],
932 finish_reason: "stop".into(),
933 usage: None,
934 })
935 }
936
937 async fn generate_with_tools(
938 &self,
939 _prompt: &str,
940 _tools: &[ToolDefinition],
941 ) -> Result<LLMResponse> {
942 Ok(LLMResponse {
943 content: format!("mock-{}", self.id),
944 tool_calls: vec![],
945 finish_reason: "stop".into(),
946 usage: None,
947 })
948 }
949
950 async fn generate_with_tools_and_history(
951 &self,
952 _messages: &[crate::coordinator::ConversationMessage],
953 _tools: &[ToolDefinition],
954 ) -> Result<LLMResponse> {
955 Ok(LLMResponse {
956 content: format!("mock-{}", self.id),
957 tool_calls: vec![],
958 finish_reason: "stop".into(),
959 usage: None,
960 })
961 }
962
963 async fn stream(
964 &self,
965 _prompt: &str,
966 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
967 Err(AppError::Internal("mock stream not implemented".into()))
968 }
969
970 async fn stream_with_system(
971 &self,
972 _system: &str,
973 _prompt: &str,
974 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
975 Err(AppError::Internal("mock stream not implemented".into()))
976 }
977
978 async fn stream_with_history(
979 &self,
980 _messages: &[(String, String)],
981 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
982 Err(AppError::Internal("mock stream not implemented".into()))
983 }
984
985 fn model_name(&self) -> &str {
986 &self.model
987 }
988 }
989}
990
991#[cfg(test)]
992mod tests {
993 use super::*;
994
995 #[test]
996 fn test_llm_response_creation() {
997 let response = LLMResponse {
998 content: "Hello".to_string(),
999 tool_calls: vec![],
1000 finish_reason: "stop".to_string(),
1001 usage: None,
1002 };
1003
1004 assert_eq!(response.content, "Hello");
1005 assert!(response.tool_calls.is_empty());
1006 assert_eq!(response.finish_reason, "stop");
1007 assert!(response.usage.is_none());
1008 }
1009
1010 #[test]
1011 fn test_llm_response_with_usage() {
1012 let usage = TokenUsage::new(100, 50);
1013 let response = LLMResponse {
1014 content: "Hello".to_string(),
1015 tool_calls: vec![],
1016 finish_reason: "stop".to_string(),
1017 usage: Some(usage),
1018 };
1019
1020 assert!(response.usage.is_some());
1021 let usage = response.usage.unwrap();
1022 assert_eq!(usage.prompt_tokens, 100);
1023 assert_eq!(usage.completion_tokens, 50);
1024 assert_eq!(usage.total_tokens, 150);
1025 }
1026
1027 #[test]
1028 fn test_llm_response_with_tool_calls() {
1029 let tool_calls = vec![
1030 ToolCall {
1031 id: "1".to_string(),
1032 name: "calculator".to_string(),
1033 arguments: serde_json::json!({"a": 1, "b": 2}),
1034 },
1035 ToolCall {
1036 id: "2".to_string(),
1037 name: "search".to_string(),
1038 arguments: serde_json::json!({"query": "test"}),
1039 },
1040 ];
1041
1042 let response = LLMResponse {
1043 content: "".to_string(),
1044 tool_calls,
1045 finish_reason: "tool_calls".to_string(),
1046 usage: Some(TokenUsage::new(50, 25)),
1047 };
1048
1049 assert_eq!(response.tool_calls.len(), 2);
1050 assert_eq!(response.tool_calls[0].name, "calculator");
1051 assert_eq!(response.finish_reason, "tool_calls");
1052 assert_eq!(response.usage.as_ref().unwrap().total_tokens, 75);
1053 }
1054
1055 #[test]
1056 fn test_factory_creation() {
1057 #[cfg(feature = "openai")]
1060 {
1061 let factory = LLMClientFactory::new(Provider::OpenAI {
1062 api_key: "sk-test".to_string(),
1063 api_base: "https://api.openai.com/v1".to_string(),
1064 model: "test".to_string(),
1065 params: ModelParams::default(),
1066 });
1067 assert_eq!(factory.default_provider().name(), "openai");
1068 }
1069 }
1070
1071 #[cfg(feature = "openai")]
1072 #[test]
1073 fn test_openai_provider_properties() {
1074 let provider = Provider::OpenAI {
1075 api_key: "sk-test".to_string(),
1076 api_base: "https://api.openai.com/v1".to_string(),
1077 model: "gpt-4".to_string(),
1078 params: ModelParams::default(),
1079 };
1080
1081 assert_eq!(provider.name(), "openai");
1082 assert!(provider.requires_api_key());
1083 assert!(!provider.is_local());
1084 }
1085
1086 #[cfg(feature = "openai")]
1087 #[test]
1088 fn test_openai_local_provider() {
1089 let provider = Provider::OpenAI {
1090 api_key: "test".to_string(),
1091 api_base: "http://localhost:8000/v1".to_string(),
1092 model: "local-model".to_string(),
1093 params: ModelParams::default(),
1094 };
1095
1096 assert!(provider.is_local());
1097 }
1098
1099 #[test]
1102 fn test_token_usage_default_all_zeros() {
1103 let usage = TokenUsage::default();
1104 assert_eq!(usage.prompt_tokens, 0);
1105 assert_eq!(usage.completion_tokens, 0);
1106 assert_eq!(usage.total_tokens, 0);
1107 }
1108
1109 #[test]
1110 fn test_token_usage_new_calculates_total() {
1111 let usage = TokenUsage::new(100, 50);
1112 assert_eq!(usage.prompt_tokens, 100);
1113 assert_eq!(usage.completion_tokens, 50);
1114 assert_eq!(usage.total_tokens, 150);
1115 }
1116
1117 #[test]
1118 fn test_token_usage_new_zero_tokens() {
1119 let usage = TokenUsage::new(0, 0);
1120 assert_eq!(usage.total_tokens, 0);
1121 }
1122
1123 #[test]
1124 fn test_token_usage_new_large_values() {
1125 let usage = TokenUsage::new(u32::MAX / 2, u32::MAX / 2 + 1);
1126 assert_eq!(usage.total_tokens, u32::MAX);
1127 }
1128
1129 #[test]
1130 fn test_token_usage_serde_roundtrip() {
1131 let usage = TokenUsage::new(100, 200);
1132 let json = serde_json::to_string(&usage).unwrap();
1133 let deserialized: TokenUsage = serde_json::from_str(&json).unwrap();
1134 assert_eq!(usage, deserialized);
1135 }
1136
1137 #[test]
1138 fn test_token_usage_serde_default_values() {
1139 let json = r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}"#;
1140 let usage: TokenUsage = serde_json::from_str(json).unwrap();
1141 assert_eq!(usage, TokenUsage::default());
1142 }
1143
1144 #[test]
1145 fn test_token_usage_serde_partial_json() {
1146 let json = r#"{"prompt_tokens":42,"completion_tokens":58,"total_tokens":100}"#;
1148 let usage: TokenUsage = serde_json::from_str(json).unwrap();
1149 assert_eq!(usage.prompt_tokens, 42);
1150 assert_eq!(usage.completion_tokens, 58);
1151 assert_eq!(usage.total_tokens, 100);
1152 }
1153
1154 #[test]
1155 fn test_token_usage_clone_eq() {
1156 let a = TokenUsage::new(10, 20);
1157 let b = a.clone();
1158 assert_eq!(a, b);
1159 }
1160
1161 #[test]
1162 fn test_token_usage_debug_format() {
1163 let usage = TokenUsage::new(1, 2);
1164 let debug_str = format!("{:?}", usage);
1165 assert!(debug_str.contains("TokenUsage"));
1166 assert!(debug_str.contains("prompt_tokens"));
1167 }
1168
1169 #[test]
1172 fn test_model_params_default_all_none() {
1173 let params = ModelParams::default();
1174 assert!(params.temperature.is_none());
1175 assert!(params.max_tokens.is_none());
1176 assert!(params.top_p.is_none());
1177 assert!(params.frequency_penalty.is_none());
1178 assert!(params.presence_penalty.is_none());
1179 }
1180
1181 #[test]
1182 fn test_model_params_from_model_config_all_fields() {
1183 let config = ModelConfig {
1184 provider: "openai".to_string(),
1185 model: "gpt-4".to_string(),
1186 temperature: 0.5,
1187 max_tokens: 1024,
1188 };
1189 let params = ModelParams::from_model_config(&config);
1190 assert_eq!(params.temperature, Some(0.5));
1191 assert_eq!(params.max_tokens, Some(1024));
1192 assert!(params.top_p.is_none());
1193 assert!(params.frequency_penalty.is_none());
1194 assert!(params.presence_penalty.is_none());
1195 }
1196
1197 #[test]
1198 fn test_model_params_from_model_config_optional_none() {
1199 let config = ModelConfig {
1200 provider: "openai".to_string(),
1201 model: "mistral".to_string(),
1202 temperature: 0.7,
1203 max_tokens: 512,
1204 };
1205 let params = ModelParams::from_model_config(&config);
1206 assert_eq!(params.temperature, Some(0.7));
1207 assert_eq!(params.max_tokens, Some(512));
1208 assert!(params.top_p.is_none());
1209 assert!(params.frequency_penalty.is_none());
1210 assert!(params.presence_penalty.is_none());
1211 }
1212
1213 #[test]
1214 fn test_model_params_clone() {
1215 let params = ModelParams {
1216 temperature: Some(0.8),
1217 max_tokens: Some(2048),
1218 top_p: Some(0.95),
1219 frequency_penalty: Some(-0.5),
1220 presence_penalty: Some(0.3),
1221 };
1222 let cloned = params.clone();
1223 assert_eq!(params.temperature, cloned.temperature);
1224 assert_eq!(params.max_tokens, cloned.max_tokens);
1225 assert_eq!(params.top_p, cloned.top_p);
1226 assert_eq!(params.frequency_penalty, cloned.frequency_penalty);
1227 assert_eq!(params.presence_penalty, cloned.presence_penalty);
1228 }
1229
1230 #[test]
1233 fn test_llm_response_empty_content() {
1234 let response = LLMResponse {
1235 content: String::new(),
1236 tool_calls: vec![],
1237 finish_reason: "stop".to_string(),
1238 usage: None,
1239 };
1240 assert!(response.content.is_empty());
1241 }
1242
1243 #[test]
1244 fn test_llm_response_clone() {
1245 let response = LLMResponse {
1246 content: "hello".to_string(),
1247 tool_calls: vec![ToolCall {
1248 id: "1".to_string(),
1249 name: "fn".to_string(),
1250 arguments: serde_json::json!({"key": "value"}),
1251 }],
1252 finish_reason: "tool_calls".to_string(),
1253 usage: Some(TokenUsage::new(10, 20)),
1254 };
1255 let cloned = response.clone();
1256 assert_eq!(cloned.content, "hello");
1257 assert_eq!(cloned.tool_calls.len(), 1);
1258 assert_eq!(cloned.tool_calls[0].name, "fn");
1259 assert_eq!(cloned.finish_reason, "tool_calls");
1260 assert_eq!(cloned.usage.unwrap().total_tokens, 30);
1261 }
1262
1263 #[cfg(feature = "openai")]
1266 mod openai_tests {
1267 use super::*;
1268
1269 #[test]
1270 fn test_openai_name() {
1271 let provider = Provider::OpenAI {
1272 api_key: "sk-test".to_string(),
1273 api_base: "https://api.openai.com/v1".to_string(),
1274 model: "gpt-4".to_string(),
1275 params: ModelParams::default(),
1276 };
1277 assert_eq!(provider.name(), "openai");
1278 }
1279
1280 #[test]
1281 fn test_openai_requires_api_key() {
1282 let provider = Provider::OpenAI {
1283 api_key: "sk-test".to_string(),
1284 api_base: "https://api.openai.com/v1".to_string(),
1285 model: "gpt-4".to_string(),
1286 params: ModelParams::default(),
1287 };
1288 assert!(provider.requires_api_key());
1289 }
1290
1291 #[test]
1292 fn test_openai_is_local_localhost() {
1293 let provider = Provider::OpenAI {
1294 api_key: "test".to_string(),
1295 api_base: "http://localhost:8000/v1".to_string(),
1296 model: "local".to_string(),
1297 params: ModelParams::default(),
1298 };
1299 assert!(provider.is_local());
1300 }
1301
1302 #[test]
1303 fn test_openai_is_local_127_0_0_1() {
1304 let provider = Provider::OpenAI {
1305 api_key: "test".to_string(),
1306 api_base: "http://127.0.0.1:8000/v1".to_string(),
1307 model: "local".to_string(),
1308 params: ModelParams::default(),
1309 };
1310 assert!(provider.is_local());
1311 }
1312
1313 #[test]
1314 fn test_openai_is_not_local_remote() {
1315 let provider = Provider::OpenAI {
1316 api_key: "sk-test".to_string(),
1317 api_base: "https://api.openai.com/v1".to_string(),
1318 model: "gpt-4".to_string(),
1319 params: ModelParams::default(),
1320 };
1321 assert!(!provider.is_local());
1322 }
1323
1324 #[test]
1325 fn test_openai_from_config_missing_env_var() {
1326 std::env::remove_var("TEST_OPENAI_MISSING_KEY");
1328 let config = ProviderConfig::OpenAI {
1329 api_key_env: "TEST_OPENAI_MISSING_KEY".to_string(),
1330 api_base: "https://api.openai.com/v1".to_string(),
1331 default_model: "gpt-4".to_string(),
1332 };
1333 let result = Provider::from_config(&config, None);
1334 assert!(result.is_err());
1335 match result.unwrap_err() {
1336 AppError::Configuration(msg) => {
1337 assert!(msg.contains("TEST_OPENAI_MISSING_KEY"));
1338 }
1339 other => panic!("Expected Configuration error, got: {:?}", other),
1340 }
1341 }
1342 }
1343
1344 #[test]
1345 fn test_token_usage_not_equal() {
1346 assert_ne!(TokenUsage::new(1, 2), TokenUsage::new(3, 4));
1347 }
1348
1349 #[test]
1350 fn test_model_params_debug_format() {
1351 let params = ModelParams::default();
1352 let debug_str = format!("{:?}", params);
1353 assert!(debug_str.contains("ModelParams"));
1354 }
1355
1356 fn test_stub_provider(model: &str) -> Provider {
1357 Provider::TestStub {
1358 model: model.to_string(),
1359 }
1360 }
1361
1362 #[test]
1363 fn test_stub_provider_properties() {
1364 let provider = test_stub_provider("unit-test");
1365 assert_eq!(provider.name(), "test-stub");
1366 assert!(!provider.requires_api_key());
1367 assert!(provider.is_local());
1368 }
1369
1370 #[tokio::test]
1371 async fn test_provider_create_client_test_stub() {
1372 let client = test_stub_provider("provider-model")
1373 .create_client()
1374 .await
1375 .expect("TestStub client");
1376 assert_eq!(client.model_name(), "provider-model");
1377 }
1378
1379 #[tokio::test]
1380 async fn test_factory_create_default_via_test_stub() {
1381 let factory = LLMClientFactory::new(test_stub_provider("factory-model"));
1382 let client = factory.create_default().await.expect("factory client");
1383 assert_eq!(client.model_name(), "factory-model");
1384 }
1385
1386 #[tokio::test]
1387 async fn test_factory_trait_create_with_provider() {
1388 let factory = LLMClientFactory::new(test_stub_provider("default"));
1389 let trait_ref: &dyn LLMClientFactoryTrait = &factory;
1390 let client = trait_ref
1391 .create_with_provider(test_stub_provider("switched"))
1392 .await
1393 .expect("switched client");
1394 assert_eq!(client.model_name(), "switched");
1395 }
1396
1397 mod llm_client_trait_tests {
1398 use super::*;
1399 use crate::client::test_support::MockLLMClient;
1400 use crate::coordinator::{ConversationMessage, MessageRole};
1401 use ares_types::types::ToolDefinition;
1402
1403 #[tokio::test]
1404 async fn test_generate_and_model_name() {
1405 let client = MockLLMClient::new("trait-model");
1406 assert_eq!(client.model_name(), "trait-model");
1407 let out = client.generate("hello").await.expect("generate");
1408 assert!(out.starts_with("mock-"));
1409 }
1410
1411 #[tokio::test]
1412 async fn test_generate_with_system() {
1413 let client = MockLLMClient::new("sys");
1414 let out = client
1415 .generate_with_system("system", "prompt")
1416 .await
1417 .expect("generate_with_system");
1418 assert!(out.starts_with("mock-"));
1419 }
1420
1421 #[tokio::test]
1422 async fn test_generate_with_history() {
1423 let client = MockLLMClient::new("hist");
1424 let messages = vec![("user".to_string(), "hi".to_string())];
1425 let response = client
1426 .generate_with_history(&messages)
1427 .await
1428 .expect("generate_with_history");
1429 assert!(response.content.starts_with("mock-"));
1430 assert_eq!(response.finish_reason, "stop");
1431 assert!(response.tool_calls.is_empty());
1432 }
1433
1434 #[tokio::test]
1435 async fn test_generate_with_tools() {
1436 let client = MockLLMClient::new("tools");
1437 let tools = vec![ToolDefinition {
1438 name: "search".to_string(),
1439 description: "Search".to_string(),
1440 parameters: serde_json::json!({"type": "object"}),
1441 }];
1442 let response = client
1443 .generate_with_tools("find docs", &tools)
1444 .await
1445 .expect("generate_with_tools");
1446 assert!(response.content.starts_with("mock-"));
1447 }
1448
1449 #[tokio::test]
1450 async fn test_generate_with_tools_and_history() {
1451 let client = MockLLMClient::new("both");
1452 let messages = vec![ConversationMessage {
1453 role: MessageRole::User,
1454 content: "run tool".to_string(),
1455 tool_calls: vec![],
1456 tool_call_id: None,
1457 }];
1458 let tools = vec![ToolDefinition {
1459 name: "calc".to_string(),
1460 description: "Calculate".to_string(),
1461 parameters: serde_json::json!({"type": "object"}),
1462 }];
1463 let response = client
1464 .generate_with_tools_and_history(&messages, &tools)
1465 .await
1466 .expect("generate_with_tools_and_history");
1467 assert!(response.content.starts_with("mock-"));
1468 }
1469
1470 #[tokio::test]
1471 async fn test_stream_methods_return_internal_error() {
1472 let client = MockLLMClient::new("stream");
1473 for result in [
1474 client.stream("hi").await,
1475 client.stream_with_system("sys", "hi").await,
1476 client
1477 .stream_with_history(&[("user".into(), "hi".into())])
1478 .await,
1479 ] {
1480 assert!(matches!(result, Err(AppError::Internal(_))));
1481 }
1482 }
1483
1484 #[test]
1485 fn default_supports_hints_is_false() {
1486 let client = MockLLMClient::new("hints");
1487 assert!(!client.supports_hints());
1488 client.set_hints(GenerationHints {
1491 json_mode: true,
1492 ..Default::default()
1493 });
1494 }
1495
1496 #[test]
1497 fn hint_recording_mock_records_set_hints_calls() {
1498 use parking_lot::Mutex;
1499 use std::sync::Arc;
1500
1501 #[derive(Default)]
1503 struct HintRecordingClient {
1504 hints: Mutex<Vec<GenerationHints>>,
1505 }
1506
1507 #[async_trait]
1508 impl LLMClient for HintRecordingClient {
1509 async fn generate(&self, _prompt: &str) -> Result<String> {
1510 Err(AppError::Internal("unused".into()))
1511 }
1512
1513 async fn generate_with_system(
1514 &self,
1515 _system: &str,
1516 _prompt: &str,
1517 ) -> Result<String> {
1518 Err(AppError::Internal("unused".into()))
1519 }
1520
1521 async fn generate_with_history(
1522 &self,
1523 _messages: &[(String, String)],
1524 ) -> Result<LLMResponse> {
1525 Err(AppError::Internal("unused".into()))
1526 }
1527
1528 async fn generate_with_tools(
1529 &self,
1530 _prompt: &str,
1531 _tools: &[ToolDefinition],
1532 ) -> Result<LLMResponse> {
1533 Err(AppError::Internal("unused".into()))
1534 }
1535
1536 async fn generate_with_tools_and_history(
1537 &self,
1538 _messages: &[crate::coordinator::ConversationMessage],
1539 _tools: &[ToolDefinition],
1540 ) -> Result<LLMResponse> {
1541 Err(AppError::Internal("unused".into()))
1542 }
1543
1544 async fn stream(
1545 &self,
1546 _prompt: &str,
1547 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1548 {
1549 Err(AppError::Internal("unused".into()))
1550 }
1551
1552 async fn stream_with_system(
1553 &self,
1554 _system: &str,
1555 _prompt: &str,
1556 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1557 {
1558 Err(AppError::Internal("unused".into()))
1559 }
1560
1561 async fn stream_with_history(
1562 &self,
1563 _messages: &[(String, String)],
1564 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>>
1565 {
1566 Err(AppError::Internal("unused".into()))
1567 }
1568
1569 fn model_name(&self) -> &str {
1570 "hint-recording-mock"
1571 }
1572
1573 fn supports_hints(&self) -> bool {
1574 true
1575 }
1576
1577 fn set_hints(&self, hints: GenerationHints) {
1578 self.hints.lock().push(hints);
1579 }
1580 }
1581
1582 let client = Arc::new(HintRecordingClient::default());
1583 assert!(client.supports_hints());
1584 client.set_hints(GenerationHints {
1585 json_mode: true,
1586 suppress_reasoning: false,
1587 max_tokens: Some(256),
1588 guided_grammar: None,
1589 });
1590 client.set_hints(GenerationHints::default());
1591
1592 let recorded = client.hints.lock();
1593 assert_eq!(
1594 recorded.len(),
1595 2,
1596 "every set_hints call is recorded in order"
1597 );
1598 assert!(recorded[0].json_mode && recorded[0].max_tokens == Some(256));
1599 assert_eq!(
1600 recorded[1],
1601 GenerationHints::default(),
1602 "clearing via Default::default() reaches the impl"
1603 );
1604 }
1605 }
1606}