1use crate::error::{AgentLoopError, Result};
18use crate::openresponses_protocol::{CompactRequest, CompactResponse};
19use crate::runtime_agent::RuntimeAgent;
20use crate::tool_types::{ToolCall, ToolDefinition};
21use async_trait::async_trait;
22use chrono::{DateTime, Utc};
23use futures::Stream;
24use std::collections::HashMap;
25use std::pin::Pin;
26use std::sync::Arc;
27
28pub type LlmResponseStream = Pin<Box<dyn Stream<Item = Result<LlmStreamEvent>> + Send>>;
34
35#[derive(Debug, Clone)]
37pub enum LlmStreamEvent {
38 TextDelta(String),
40 ThinkingDelta(String),
42 ThinkingSignature(String),
45 ReasonItem {
51 provider: String,
53 model: Option<String>,
55 item_id: String,
57 encrypted_content: Option<String>,
59 summary: Vec<String>,
61 token_count: Option<u32>,
63 },
64 ToolCalls(Vec<ToolCall>),
66 Done(Box<LlmCompletionMetadata>),
68 Error(String),
70}
71
72#[derive(Debug, Clone)]
83pub struct DiscoveredModel {
84 pub model_id: String,
86 pub display_name: Option<String>,
88 pub created_at: Option<DateTime<Utc>>,
90 pub owned_by: Option<String>,
92 pub discovered_profile: Option<crate::llm_models::LlmModelProfile>,
95}
96
97#[derive(Debug, Clone, Default)]
105pub struct LlmCompletionMetadata {
106 pub total_tokens: Option<u32>,
108 pub prompt_tokens: Option<u32>,
110 pub completion_tokens: Option<u32>,
112 pub cache_read_tokens: Option<u32>,
114 pub cache_creation_tokens: Option<u32>,
116 pub provider_cost_usd: Option<f64>,
120 pub model: Option<String>,
122 pub finish_reason: Option<String>,
124 pub retry_metadata: Option<crate::llm_retry::RetryMetadata>,
126 pub response_id: Option<String>,
129 pub phase: Option<String>,
133}
134
135#[async_trait]
139pub trait LlmDriver: Send + Sync {
140 async fn chat_completion_stream(
142 &self,
143 messages: Vec<LlmMessage>,
144 config: &LlmCallConfig,
145 ) -> Result<LlmResponseStream>;
146
147 async fn chat_completion(
149 &self,
150 messages: Vec<LlmMessage>,
151 config: &LlmCallConfig,
152 ) -> Result<LlmResponse> {
153 use futures::StreamExt;
154
155 let mut stream = self.chat_completion_stream(messages, config).await?;
156 let mut text = String::new();
157 let mut thinking = String::new();
158 let mut thinking_signature: Option<String> = None;
159 let mut tool_calls = Vec::new();
160 let mut metadata = LlmCompletionMetadata::default();
161
162 while let Some(event) = stream.next().await {
163 match event? {
164 LlmStreamEvent::TextDelta(delta) => text.push_str(&delta),
165 LlmStreamEvent::ThinkingDelta(delta) => thinking.push_str(&delta),
166 LlmStreamEvent::ThinkingSignature(sig) => thinking_signature = Some(sig),
167 LlmStreamEvent::ReasonItem {
168 encrypted_content, ..
169 } => {
170 if let Some(sig) = encrypted_content {
171 thinking_signature = Some(sig);
172 }
173 }
174 LlmStreamEvent::ToolCalls(calls) => tool_calls = calls,
175 LlmStreamEvent::Done(meta) => metadata = *meta,
176 LlmStreamEvent::Error(err) => return Err(crate::error::AgentLoopError::llm(err)),
177 }
178 }
179
180 Ok(LlmResponse {
181 text,
182 thinking: if thinking.is_empty() {
183 None
184 } else {
185 Some(thinking)
186 },
187 thinking_signature,
188 tool_calls: if tool_calls.is_empty() {
189 None
190 } else {
191 Some(tool_calls)
192 },
193 metadata,
194 })
195 }
196
197 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
205 Ok(None)
207 }
208
209 fn supports_compact(&self) -> bool {
218 false
220 }
221
222 async fn compact(&self, _request: CompactRequest) -> Result<Option<CompactResponse>> {
242 Ok(None)
244 }
245}
246
247#[async_trait]
249impl LlmDriver for Box<dyn LlmDriver> {
250 async fn chat_completion_stream(
251 &self,
252 messages: Vec<LlmMessage>,
253 config: &LlmCallConfig,
254 ) -> Result<LlmResponseStream> {
255 (**self).chat_completion_stream(messages, config).await
256 }
257
258 async fn chat_completion(
259 &self,
260 messages: Vec<LlmMessage>,
261 config: &LlmCallConfig,
262 ) -> Result<LlmResponse> {
263 (**self).chat_completion(messages, config).await
264 }
265
266 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
267 (**self).list_models().await
268 }
269
270 fn supports_compact(&self) -> bool {
271 (**self).supports_compact()
272 }
273
274 async fn compact(&self, request: CompactRequest) -> Result<Option<CompactResponse>> {
275 (**self).compact(request).await
276 }
277}
278
279#[derive(Debug, Clone)]
285pub struct LlmMessage {
286 pub role: LlmMessageRole,
287 pub content: LlmMessageContent,
288 pub tool_calls: Option<Vec<ToolCall>>,
289 pub tool_call_id: Option<String>,
290 pub phase: Option<crate::message::ExecutionPhase>,
295 pub thinking: Option<String>,
298 pub thinking_signature: Option<String>,
301}
302
303impl LlmMessage {
304 pub fn text(role: LlmMessageRole, content: impl Into<String>) -> Self {
306 Self {
307 role,
308 content: LlmMessageContent::Text(content.into()),
309 tool_calls: None,
310 tool_call_id: None,
311 phase: None,
312 thinking: None,
313 thinking_signature: None,
314 }
315 }
316
317 pub fn parts(role: LlmMessageRole, parts: Vec<LlmContentPart>) -> Self {
319 Self {
320 role,
321 content: LlmMessageContent::Parts(parts),
322 tool_calls: None,
323 tool_call_id: None,
324 phase: None,
325 thinking: None,
326 thinking_signature: None,
327 }
328 }
329
330 pub fn content_as_text(&self) -> String {
332 self.content.to_text()
333 }
334
335 pub fn prepend_text_prefix(&mut self, prefix: &str) {
340 match &mut self.content {
341 LlmMessageContent::Text(text) => {
342 *text = format!("{}{}", prefix, text);
343 }
344 LlmMessageContent::Parts(parts) => {
345 for part in parts.iter_mut() {
346 if let LlmContentPart::Text { text } = part {
347 *text = format!("{}{}", prefix, text);
348 return;
349 }
350 }
351 parts.insert(
353 0,
354 LlmContentPart::Text {
355 text: prefix.to_string(),
356 },
357 );
358 }
359 }
360 }
361}
362
363#[derive(Debug, Clone)]
365pub enum LlmMessageContent {
366 Text(String),
368 Parts(Vec<LlmContentPart>),
370}
371
372impl LlmMessageContent {
373 pub fn to_text(&self) -> String {
375 match self {
376 LlmMessageContent::Text(s) => s.clone(),
377 LlmMessageContent::Parts(parts) => parts
378 .iter()
379 .filter_map(|p| match p {
380 LlmContentPart::Text { text } => Some(text.clone()),
381 _ => None,
382 })
383 .collect::<Vec<_>>()
384 .join(""),
385 }
386 }
387
388 pub fn is_text(&self) -> bool {
390 matches!(self, LlmMessageContent::Text(_))
391 }
392
393 pub fn is_parts(&self) -> bool {
395 matches!(self, LlmMessageContent::Parts(_))
396 }
397}
398
399impl From<String> for LlmMessageContent {
400 fn from(s: String) -> Self {
401 LlmMessageContent::Text(s)
402 }
403}
404
405impl From<&str> for LlmMessageContent {
406 fn from(s: &str) -> Self {
407 LlmMessageContent::Text(s.to_string())
408 }
409}
410
411#[derive(Debug, Clone)]
413pub enum LlmContentPart {
414 Text { text: String },
416 Image { url: String },
418 Audio { url: String },
420}
421
422impl LlmContentPart {
423 pub fn text(text: impl Into<String>) -> Self {
425 LlmContentPart::Text { text: text.into() }
426 }
427
428 pub fn image(url: impl Into<String>) -> Self {
430 LlmContentPart::Image { url: url.into() }
431 }
432
433 pub fn audio(url: impl Into<String>) -> Self {
435 LlmContentPart::Audio { url: url.into() }
436 }
437}
438
439#[derive(Debug, Clone, PartialEq, Eq)]
441pub enum LlmMessageRole {
442 System,
443 User,
444 Assistant,
445 Tool,
446}
447
448#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
458pub struct ToolSearchConfig {
459 pub enabled: bool,
461 pub threshold: usize,
464}
465
466#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
468#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
469#[serde(rename_all = "snake_case")]
470pub enum PromptCacheStrategy {
471 #[default]
473 Auto,
474}
475
476#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
481#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
482pub struct PromptCacheConfig {
483 pub enabled: bool,
485 #[serde(default)]
487 pub strategy: PromptCacheStrategy,
488 #[serde(default, skip_serializing_if = "Option::is_none")]
495 pub gemini_cached_content: Option<String>,
496}
497
498#[derive(Debug, Clone)]
500pub struct LlmCallConfig {
501 pub model: String,
502 pub temperature: Option<f32>,
503 pub max_tokens: Option<u32>,
504 pub tools: Vec<ToolDefinition>,
505 pub reasoning_effort: Option<String>,
507 pub metadata: HashMap<String, String>,
511 pub previous_response_id: Option<String>,
514 pub tool_search: Option<ToolSearchConfig>,
516 pub prompt_cache: Option<PromptCacheConfig>,
518}
519
520impl From<&RuntimeAgent> for LlmCallConfig {
521 fn from(runtime_agent: &RuntimeAgent) -> Self {
522 Self {
523 model: runtime_agent.model.clone(),
524 temperature: runtime_agent.temperature,
525 max_tokens: runtime_agent.max_tokens,
526 tools: runtime_agent.tools.clone(),
527 reasoning_effort: None, metadata: HashMap::new(), previous_response_id: None,
530 tool_search: runtime_agent.tool_search.clone(),
531 prompt_cache: runtime_agent.prompt_cache.clone(),
532 }
533 }
534}
535
536#[derive(Debug, Clone)]
538pub struct LlmResponse {
539 pub text: String,
540 pub thinking: Option<String>,
542 pub thinking_signature: Option<String>,
544 pub tool_calls: Option<Vec<ToolCall>>,
545 pub metadata: LlmCompletionMetadata,
546}
547
548pub struct LlmCallConfigBuilder {
567 config: LlmCallConfig,
568}
569
570impl LlmCallConfigBuilder {
571 pub fn from(runtime_agent: &RuntimeAgent) -> Self {
573 Self {
574 config: LlmCallConfig::from(runtime_agent),
575 }
576 }
577
578 pub fn reasoning_effort(mut self, effort: impl Into<String>) -> Self {
580 self.config.reasoning_effort = Some(effort.into());
581 self
582 }
583
584 pub fn model(mut self, model: impl Into<String>) -> Self {
586 self.config.model = model.into();
587 self
588 }
589
590 pub fn temperature(mut self, temp: f32) -> Self {
592 self.config.temperature = Some(temp);
593 self
594 }
595
596 pub fn max_tokens(mut self, tokens: u32) -> Self {
598 self.config.max_tokens = Some(tokens);
599 self
600 }
601
602 pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
604 self.config.tools = tools;
605 self
606 }
607
608 pub fn metadata(mut self, metadata: HashMap<String, String>) -> Self {
613 self.config.metadata = metadata;
614 self
615 }
616
617 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
619 self.config.metadata.insert(key.into(), value.into());
620 self
621 }
622
623 pub fn previous_response_id(mut self, id: Option<String>) -> Self {
625 self.config.previous_response_id = id;
626 self
627 }
628
629 pub fn tool_search(mut self, config: ToolSearchConfig) -> Self {
631 self.config.tool_search = Some(config);
632 self
633 }
634
635 pub fn prompt_cache(mut self, config: PromptCacheConfig) -> Self {
637 self.config.prompt_cache = Some(config);
638 self
639 }
640
641 pub fn build(self) -> LlmCallConfig {
643 self.config
644 }
645}
646
647impl From<&crate::message::Message> for LlmMessage {
652 fn from(msg: &crate::message::Message) -> Self {
658 let role = match msg.role {
659 crate::message::MessageRole::System => LlmMessageRole::System,
660 crate::message::MessageRole::User => LlmMessageRole::User,
661 crate::message::MessageRole::Agent => LlmMessageRole::Assistant,
662 crate::message::MessageRole::ToolResult => LlmMessageRole::Tool,
663 };
664
665 let tool_calls: Vec<ToolCall> = msg
667 .tool_calls()
668 .into_iter()
669 .map(|tc| ToolCall {
670 id: tc.id.clone(),
671 name: tc.name.clone(),
672 arguments: tc.arguments.clone(),
673 })
674 .collect();
675
676 LlmMessage {
677 role,
678 content: LlmMessageContent::Text(msg.content_to_llm_string()),
679 tool_calls: if tool_calls.is_empty() {
680 None
681 } else {
682 Some(tool_calls)
683 },
684 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
685 phase: msg.phase,
686 thinking: msg.thinking.clone(),
687 thinking_signature: msg.thinking_signature.clone(),
688 }
689 }
690}
691
692use crate::traits::ResolvedImage;
697use uuid::Uuid;
698
699impl LlmMessage {
700 pub fn from_message_with_images(
720 msg: &crate::message::Message,
721 resolved_images: &HashMap<Uuid, ResolvedImage>,
722 ) -> Self {
723 use crate::message::{ContentPart, MessageRole};
724
725 let role = match msg.role {
726 MessageRole::System => LlmMessageRole::System,
727 MessageRole::User => LlmMessageRole::User,
728 MessageRole::Agent => LlmMessageRole::Assistant,
729 MessageRole::ToolResult => LlmMessageRole::Tool,
730 };
731
732 let mut parts: Vec<LlmContentPart> = Vec::new();
734 let mut tool_calls: Vec<ToolCall> = Vec::new();
735
736 for part in &msg.content {
737 match part {
738 ContentPart::Text(t) => {
739 parts.push(LlmContentPart::Text {
740 text: t.text.clone(),
741 });
742 }
743 ContentPart::Image(img) => {
744 if let Some(url) = &img.url {
746 parts.push(LlmContentPart::Image { url: url.clone() });
747 } else if let (Some(base64), Some(media_type)) = (&img.base64, &img.media_type)
748 {
749 let data_url = format!("data:{};base64,{}", media_type, base64);
750 parts.push(LlmContentPart::Image { url: data_url });
751 }
752 }
753 ContentPart::ImageFile(img_file) => {
754 if let Some(resolved) = resolved_images.get(&img_file.image_id.uuid()) {
756 parts.push(LlmContentPart::Image {
757 url: resolved.to_data_url(),
758 });
759 } else {
760 parts.push(LlmContentPart::Text {
762 text: format!("[Image not found: {}]", img_file.image_id),
763 });
764 }
765 }
766 ContentPart::ToolCall(tc) => {
767 tool_calls.push(ToolCall {
769 id: tc.id.clone(),
770 name: tc.name.clone(),
771 arguments: tc.arguments.clone(),
772 });
773 }
774 ContentPart::ToolResult(tr) => {
775 let text = if let Some(err) = &tr.error {
777 format!("Tool error: {}", err)
778 } else if let Some(res) = &tr.result {
779 serde_json::to_string(res).unwrap_or_else(|_| "{}".to_string())
780 } else {
781 "{}".to_string()
782 };
783 let text = truncate_tool_result(text);
787 parts.push(LlmContentPart::Text { text });
788 }
789 }
790 }
791
792 let content = if parts.len() == 1 && matches!(&parts[0], LlmContentPart::Text { .. }) {
794 if let LlmContentPart::Text { text } = &parts[0] {
796 LlmMessageContent::Text(text.clone())
797 } else {
798 LlmMessageContent::Parts(parts)
799 }
800 } else if parts.is_empty() {
801 LlmMessageContent::Text(String::new())
803 } else {
804 LlmMessageContent::Parts(parts)
806 };
807
808 LlmMessage {
809 role,
810 content,
811 tool_calls: if tool_calls.is_empty() {
812 None
813 } else {
814 Some(tool_calls)
815 },
816 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
817 phase: msg.phase,
818 thinking: msg.thinking.clone(),
819 thinking_signature: msg.thinking_signature.clone(),
820 }
821 }
822
823 pub fn message_has_image_files(msg: &crate::message::Message) -> bool {
825 msg.content.iter().any(|p| p.is_image_file())
826 }
827
828 pub fn extract_image_file_ids(msg: &crate::message::Message) -> Vec<Uuid> {
830 msg.content
831 .iter()
832 .filter_map(|p| match p {
833 crate::message::ContentPart::ImageFile(f) => Some(f.image_id.uuid()),
834 _ => None,
835 })
836 .collect()
837 }
838}
839
840#[derive(Debug, Clone, PartialEq, Eq, Hash)]
846pub enum ProviderType {
847 OpenAI,
850 AzureOpenAI,
852 OpenAICompletions,
855 Anthropic,
856 Gemini,
858 LlmSim,
860}
861
862impl std::str::FromStr for ProviderType {
863 type Err = String;
864
865 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
866 match s.to_lowercase().as_str() {
867 "openai" => Ok(ProviderType::OpenAI),
868 "azure_openai" => Ok(ProviderType::AzureOpenAI),
869 "openai_completions" => Ok(ProviderType::OpenAICompletions),
870 "anthropic" => Ok(ProviderType::Anthropic),
871 "gemini" => Ok(ProviderType::Gemini),
872 "llmsim" => Ok(ProviderType::LlmSim),
873 _ => Err(format!("Unknown provider type: {}", s)),
874 }
875 }
876}
877
878impl std::fmt::Display for ProviderType {
879 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
880 match self {
881 ProviderType::OpenAI => write!(f, "openai"),
882 ProviderType::AzureOpenAI => write!(f, "azure_openai"),
883 ProviderType::OpenAICompletions => write!(f, "openai_completions"),
884 ProviderType::Anthropic => write!(f, "anthropic"),
885 ProviderType::Gemini => write!(f, "gemini"),
886 ProviderType::LlmSim => write!(f, "llmsim"),
887 }
888 }
889}
890
891#[derive(Debug, Clone)]
893pub struct ProviderConfig {
894 pub provider_type: ProviderType,
896 pub api_key: Option<String>,
898 pub base_url: Option<String>,
900}
901
902impl ProviderConfig {
903 pub fn new(provider_type: ProviderType) -> Self {
905 Self {
906 provider_type,
907 api_key: None,
908 base_url: None,
909 }
910 }
911
912 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
914 self.api_key = Some(api_key.into());
915 self
916 }
917
918 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
920 self.base_url = Some(base_url.into());
921 self
922 }
923}
924
925pub type BoxedLlmDriver = Box<dyn LlmDriver>;
927
928pub type DriverFactory = Arc<dyn Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync>;
936
937#[derive(Clone, Default)]
957pub struct DriverRegistry {
958 factories: HashMap<ProviderType, DriverFactory>,
959}
960
961impl DriverRegistry {
962 pub fn new() -> Self {
964 Self {
965 factories: HashMap::new(),
966 }
967 }
968
969 pub fn register<F>(&mut self, provider_type: ProviderType, factory: F)
971 where
972 F: Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync + 'static,
973 {
974 self.factories.insert(provider_type, Arc::new(factory));
975 }
976
977 pub fn create_driver(&self, config: &ProviderConfig) -> Result<BoxedLlmDriver> {
985 let api_key = if config.provider_type == ProviderType::LlmSim {
987 config.api_key.as_deref().unwrap_or("")
989 } else {
990 config.api_key.as_ref().ok_or_else(|| {
991 AgentLoopError::llm(
992 "API key is required. Configure the API key in provider settings.",
993 )
994 })?
995 };
996
997 let factory = self.factories.get(&config.provider_type).ok_or_else(|| {
999 AgentLoopError::driver_not_registered(config.provider_type.to_string())
1000 })?;
1001
1002 Ok(factory(api_key, config.base_url.as_deref()))
1004 }
1005
1006 pub fn has_driver(&self, provider_type: &ProviderType) -> bool {
1008 self.factories.contains_key(provider_type)
1009 }
1010
1011 pub fn registered_providers(&self) -> Vec<ProviderType> {
1013 self.factories.keys().cloned().collect()
1014 }
1015}
1016
1017const MAX_TOOL_RESULT_BYTES: usize = 64 * 1024;
1022
1023const TRUNCATION_SUFFIX: &str =
1024 "\n\n[Output truncated — exceeded 64 KiB limit. Try quiet flags, pipes, or redirect to file.]";
1025
1026fn truncate_tool_result(text: String) -> String {
1027 if text.len() <= MAX_TOOL_RESULT_BYTES {
1028 return text;
1029 }
1030 let content_budget = MAX_TOOL_RESULT_BYTES.saturating_sub(TRUNCATION_SUFFIX.len());
1031 let mut end = content_budget;
1032 while end > 0 && !text.is_char_boundary(end) {
1033 end -= 1;
1034 }
1035 let mut truncated = text[..end].to_string();
1036 truncated.push_str(TRUNCATION_SUFFIX);
1037 truncated
1038}
1039
1040#[cfg(test)]
1045mod tests {
1046 use super::*;
1047
1048 #[test]
1049 fn test_llm_call_config_builder_from_runtime_agent() {
1050 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1051 let llm_config = LlmCallConfigBuilder::from(&runtime_agent).build();
1052
1053 assert_eq!(llm_config.model, "gpt-4o");
1054 assert!(llm_config.reasoning_effort.is_none());
1055 assert!(llm_config.temperature.is_none());
1056 assert!(llm_config.max_tokens.is_none());
1057 assert!(llm_config.tools.is_empty());
1058 assert!(llm_config.metadata.is_empty());
1059 }
1060
1061 #[test]
1062 fn test_llm_call_config_builder_with_metadata() {
1063 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1064 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1065 .with_metadata("session_id", "session_abc123")
1066 .with_metadata("agent_id", "agent_xyz789")
1067 .build();
1068
1069 assert_eq!(
1070 llm_config.metadata.get("session_id"),
1071 Some(&"session_abc123".to_string())
1072 );
1073 assert_eq!(
1074 llm_config.metadata.get("agent_id"),
1075 Some(&"agent_xyz789".to_string())
1076 );
1077 }
1078
1079 #[test]
1080 fn test_llm_call_config_builder_with_metadata_hashmap() {
1081 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1082 let mut metadata = HashMap::new();
1083 metadata.insert("key1".to_string(), "value1".to_string());
1084 metadata.insert("key2".to_string(), "value2".to_string());
1085
1086 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1087 .metadata(metadata)
1088 .build();
1089
1090 assert_eq!(llm_config.metadata.get("key1"), Some(&"value1".to_string()));
1091 assert_eq!(llm_config.metadata.get("key2"), Some(&"value2".to_string()));
1092 }
1093
1094 #[test]
1095 fn test_llm_call_config_builder_with_reasoning_effort() {
1096 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1097 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1098 .reasoning_effort("high")
1099 .build();
1100
1101 assert_eq!(llm_config.reasoning_effort, Some("high".to_string()));
1102 }
1103
1104 #[test]
1105 fn test_llm_call_config_builder_with_all_options() {
1106 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1107 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1108 .model("claude-3-opus")
1109 .reasoning_effort("medium")
1110 .temperature(0.7)
1111 .max_tokens(1000)
1112 .build();
1113
1114 assert_eq!(llm_config.model, "claude-3-opus");
1115 assert_eq!(llm_config.reasoning_effort, Some("medium".to_string()));
1116 assert_eq!(llm_config.temperature, Some(0.7));
1117 assert_eq!(llm_config.max_tokens, Some(1000));
1118 }
1119
1120 #[test]
1121 fn test_provider_type_parsing() {
1122 assert_eq!(
1123 "openai".parse::<ProviderType>().unwrap(),
1124 ProviderType::OpenAI
1125 );
1126 assert_eq!(
1127 "openai_completions".parse::<ProviderType>().unwrap(),
1128 ProviderType::OpenAICompletions
1129 );
1130 assert_eq!(
1131 "azure_openai".parse::<ProviderType>().unwrap(),
1132 ProviderType::AzureOpenAI
1133 );
1134 assert_eq!(
1135 "anthropic".parse::<ProviderType>().unwrap(),
1136 ProviderType::Anthropic
1137 );
1138 assert_eq!(
1139 "gemini".parse::<ProviderType>().unwrap(),
1140 ProviderType::Gemini
1141 );
1142 assert!("ollama".parse::<ProviderType>().is_err());
1144 assert!("custom".parse::<ProviderType>().is_err());
1145 }
1146
1147 #[test]
1148 fn test_provider_type_display() {
1149 assert_eq!(ProviderType::OpenAI.to_string(), "openai");
1150 assert_eq!(ProviderType::AzureOpenAI.to_string(), "azure_openai");
1151 assert_eq!(
1152 ProviderType::OpenAICompletions.to_string(),
1153 "openai_completions"
1154 );
1155 assert_eq!(ProviderType::Anthropic.to_string(), "anthropic");
1156 assert_eq!(ProviderType::Gemini.to_string(), "gemini");
1157 }
1158
1159 #[test]
1160 fn test_provider_config_builder() {
1161 let config = ProviderConfig::new(ProviderType::Anthropic)
1162 .with_api_key("test-key")
1163 .with_base_url("https://custom.api.com");
1164
1165 assert_eq!(config.provider_type, ProviderType::Anthropic);
1166 assert_eq!(config.api_key, Some("test-key".to_string()));
1167 assert_eq!(config.base_url, Some("https://custom.api.com".to_string()));
1168 }
1169
1170 #[test]
1171 fn test_driver_registry_requires_api_key() {
1172 let mut registry = DriverRegistry::new();
1174 registry.register(ProviderType::OpenAI, |_api_key, _base_url| {
1175 struct MockDriver;
1177 #[async_trait]
1178 impl LlmDriver for MockDriver {
1179 async fn chat_completion_stream(
1180 &self,
1181 _messages: Vec<LlmMessage>,
1182 _config: &LlmCallConfig,
1183 ) -> Result<LlmResponseStream> {
1184 unimplemented!()
1185 }
1186 }
1187 Box::new(MockDriver)
1188 });
1189
1190 let config = ProviderConfig::new(ProviderType::OpenAI);
1192 let result = registry.create_driver(&config);
1193 assert!(result.is_err());
1194
1195 let config_with_key = ProviderConfig::new(ProviderType::OpenAI).with_api_key("test-key");
1197 let result = registry.create_driver(&config_with_key);
1198 assert!(result.is_ok());
1199 }
1200
1201 #[test]
1202 fn test_driver_registry_returns_error_for_unregistered_provider() {
1203 let registry = DriverRegistry::new();
1204 let config = ProviderConfig::new(ProviderType::Anthropic).with_api_key("test-key");
1205
1206 let result = registry.create_driver(&config);
1207
1208 if let Err(AgentLoopError::DriverNotRegistered(provider)) = result {
1210 assert_eq!(provider, "anthropic");
1211 } else {
1212 panic!("Expected DriverNotRegistered error");
1213 }
1214 }
1215
1216 #[test]
1217 fn test_driver_registry_registration() {
1218 let mut registry = DriverRegistry::new();
1219
1220 assert!(!registry.has_driver(&ProviderType::OpenAI));
1221 assert!(!registry.has_driver(&ProviderType::Anthropic));
1222
1223 registry.register(ProviderType::OpenAI, |_, _| {
1224 struct MockDriver;
1225 #[async_trait]
1226 impl LlmDriver for MockDriver {
1227 async fn chat_completion_stream(
1228 &self,
1229 _messages: Vec<LlmMessage>,
1230 _config: &LlmCallConfig,
1231 ) -> Result<LlmResponseStream> {
1232 unimplemented!()
1233 }
1234 }
1235 Box::new(MockDriver)
1236 });
1237
1238 assert!(registry.has_driver(&ProviderType::OpenAI));
1239 assert!(!registry.has_driver(&ProviderType::Anthropic));
1240 }
1241
1242 use crate::{ContentPart, ImageFileContentPart, Message, MessageRole, TextContentPart};
1247
1248 #[test]
1249 fn test_message_has_image_files_with_image_file() {
1250 let message = Message {
1251 id: uuid::Uuid::new_v4().into(),
1252 role: MessageRole::User,
1253 content: vec![
1254 ContentPart::Text(TextContentPart {
1255 text: "Look at this image".to_string(),
1256 }),
1257 ContentPart::ImageFile(ImageFileContentPart {
1258 image_id: uuid::Uuid::new_v4().into(),
1259 filename: Some("test.png".to_string()),
1260 }),
1261 ],
1262 phase: None,
1263 thinking: None,
1264 thinking_signature: None,
1265 controls: None,
1266 metadata: None,
1267 external_actor: None,
1268 created_at: chrono::Utc::now(),
1269 };
1270
1271 assert!(LlmMessage::message_has_image_files(&message));
1272 }
1273
1274 #[test]
1275 fn test_message_has_image_files_without_image_file() {
1276 let message = Message {
1277 id: uuid::Uuid::new_v4().into(),
1278 role: MessageRole::User,
1279 content: vec![ContentPart::Text(TextContentPart {
1280 text: "Just text".to_string(),
1281 })],
1282 phase: None,
1283 thinking: None,
1284 thinking_signature: None,
1285 controls: None,
1286 metadata: None,
1287 external_actor: None,
1288 created_at: chrono::Utc::now(),
1289 };
1290
1291 assert!(!LlmMessage::message_has_image_files(&message));
1292 }
1293
1294 #[test]
1295 fn test_extract_image_file_ids() {
1296 let id1 = uuid::Uuid::new_v4();
1297 let id2 = uuid::Uuid::new_v4();
1298
1299 let message = Message {
1300 id: uuid::Uuid::new_v4().into(),
1301 role: MessageRole::User,
1302 content: vec![
1303 ContentPart::Text(TextContentPart {
1304 text: "Look at these images".to_string(),
1305 }),
1306 ContentPart::ImageFile(ImageFileContentPart {
1307 image_id: id1.into(),
1308 filename: Some("test1.png".to_string()),
1309 }),
1310 ContentPart::ImageFile(ImageFileContentPart {
1311 image_id: id2.into(),
1312 filename: Some("test2.png".to_string()),
1313 }),
1314 ],
1315 phase: None,
1316 thinking: None,
1317 thinking_signature: None,
1318 controls: None,
1319 metadata: None,
1320 external_actor: None,
1321 created_at: chrono::Utc::now(),
1322 };
1323
1324 let ids = LlmMessage::extract_image_file_ids(&message);
1325 assert_eq!(ids.len(), 2);
1326 assert!(ids.contains(&id1));
1327 assert!(ids.contains(&id2));
1328 }
1329
1330 #[test]
1331 fn test_from_message_with_images_text_only() {
1332 let message = Message {
1333 id: uuid::Uuid::new_v4().into(),
1334 role: MessageRole::User,
1335 content: vec![ContentPart::Text(TextContentPart {
1336 text: "Hello".to_string(),
1337 })],
1338 phase: None,
1339 thinking: None,
1340 thinking_signature: None,
1341 controls: None,
1342 metadata: None,
1343 external_actor: None,
1344 created_at: chrono::Utc::now(),
1345 };
1346
1347 let resolved = std::collections::HashMap::new();
1348 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1349
1350 assert_eq!(llm_message.role, LlmMessageRole::User);
1351 match llm_message.content {
1352 LlmMessageContent::Text(text) => assert_eq!(text, "Hello"),
1353 _ => panic!("Expected text content"),
1354 }
1355 }
1356
1357 #[test]
1358 fn test_from_message_with_images_resolved_image() {
1359 let image_id = uuid::Uuid::new_v4();
1360 let message = Message {
1361 id: uuid::Uuid::new_v4().into(),
1362 role: MessageRole::User,
1363 content: vec![
1364 ContentPart::Text(TextContentPart {
1365 text: "Look at this".to_string(),
1366 }),
1367 ContentPart::ImageFile(ImageFileContentPart {
1368 image_id: image_id.into(),
1369 filename: Some("test.png".to_string()),
1370 }),
1371 ],
1372 phase: None,
1373 thinking: None,
1374 thinking_signature: None,
1375 controls: None,
1376 metadata: None,
1377 external_actor: None,
1378 created_at: chrono::Utc::now(),
1379 };
1380
1381 let mut resolved = std::collections::HashMap::new();
1382 resolved.insert(
1383 image_id,
1384 crate::ResolvedImage::new("base64data", "image/png"),
1385 );
1386
1387 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1388
1389 match &llm_message.content {
1390 LlmMessageContent::Parts(parts) => {
1391 assert_eq!(parts.len(), 2);
1392 assert!(matches!(&parts[0], LlmContentPart::Text { .. }));
1394 if let LlmContentPart::Image { url } = &parts[1] {
1396 assert!(url.starts_with("data:image/png;base64,"));
1397 } else {
1398 panic!("Expected image content part");
1399 }
1400 }
1401 _ => panic!("Expected parts content"),
1402 }
1403 }
1404
1405 #[test]
1406 fn test_from_message_with_images_unresolved_image() {
1407 let image_id = uuid::Uuid::new_v4();
1408 let message = Message {
1409 id: uuid::Uuid::new_v4().into(),
1410 role: MessageRole::User,
1411 content: vec![ContentPart::ImageFile(ImageFileContentPart {
1412 image_id: image_id.into(),
1413 filename: Some("missing.png".to_string()),
1414 })],
1415 phase: None,
1416 thinking: None,
1417 thinking_signature: None,
1418 controls: None,
1419 metadata: None,
1420 external_actor: None,
1421 created_at: chrono::Utc::now(),
1422 };
1423
1424 let resolved = std::collections::HashMap::new();
1426 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1427
1428 match &llm_message.content {
1431 LlmMessageContent::Text(text) => {
1432 assert!(text.contains("Image not found"));
1433 }
1434 LlmMessageContent::Parts(parts) => {
1435 assert_eq!(parts.len(), 1);
1436 if let LlmContentPart::Text { text } = &parts[0] {
1437 assert!(text.contains("Image not found"));
1438 } else {
1439 panic!("Expected text placeholder for missing image");
1440 }
1441 }
1442 }
1443 }
1444
1445 #[test]
1446 fn test_prepend_text_prefix_simple_text() {
1447 let mut msg = LlmMessage::text(LlmMessageRole::User, "Hello bot");
1448 msg.prepend_text_prefix("[Alice] ");
1449 assert_eq!(msg.content_as_text(), "[Alice] Hello bot");
1450 }
1451
1452 #[test]
1453 fn test_prepend_text_prefix_parts() {
1454 let mut msg = LlmMessage::parts(
1455 LlmMessageRole::User,
1456 vec![
1457 LlmContentPart::Text {
1458 text: "Hello".to_string(),
1459 },
1460 LlmContentPart::Image {
1461 url: "data:image/png;base64,abc".to_string(),
1462 },
1463 ],
1464 );
1465 msg.prepend_text_prefix("[Bob] ");
1466 match &msg.content {
1467 LlmMessageContent::Parts(parts) => {
1468 if let LlmContentPart::Text { text } = &parts[0] {
1469 assert_eq!(text, "[Bob] Hello");
1470 } else {
1471 panic!("Expected text part");
1472 }
1473 }
1474 _ => panic!("Expected parts content"),
1475 }
1476 }
1477
1478 #[test]
1479 fn test_prepend_text_prefix_parts_no_text() {
1480 let mut msg = LlmMessage::parts(
1481 LlmMessageRole::User,
1482 vec![LlmContentPart::Image {
1483 url: "data:image/png;base64,abc".to_string(),
1484 }],
1485 );
1486 msg.prepend_text_prefix("[Eve] ");
1487 match &msg.content {
1488 LlmMessageContent::Parts(parts) => {
1489 assert_eq!(parts.len(), 2);
1490 if let LlmContentPart::Text { text } = &parts[0] {
1491 assert_eq!(text, "[Eve] ");
1492 } else {
1493 panic!("Expected prepended text part");
1494 }
1495 }
1496 _ => panic!("Expected parts content"),
1497 }
1498 }
1499}