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 model: Option<String>,
118 pub finish_reason: Option<String>,
120 pub retry_metadata: Option<crate::llm_retry::RetryMetadata>,
122 pub response_id: Option<String>,
125 pub phase: Option<String>,
129}
130
131#[async_trait]
135pub trait LlmDriver: Send + Sync {
136 async fn chat_completion_stream(
138 &self,
139 messages: Vec<LlmMessage>,
140 config: &LlmCallConfig,
141 ) -> Result<LlmResponseStream>;
142
143 async fn chat_completion(
145 &self,
146 messages: Vec<LlmMessage>,
147 config: &LlmCallConfig,
148 ) -> Result<LlmResponse> {
149 use futures::StreamExt;
150
151 let mut stream = self.chat_completion_stream(messages, config).await?;
152 let mut text = String::new();
153 let mut thinking = String::new();
154 let mut thinking_signature: Option<String> = None;
155 let mut tool_calls = Vec::new();
156 let mut metadata = LlmCompletionMetadata::default();
157
158 while let Some(event) = stream.next().await {
159 match event? {
160 LlmStreamEvent::TextDelta(delta) => text.push_str(&delta),
161 LlmStreamEvent::ThinkingDelta(delta) => thinking.push_str(&delta),
162 LlmStreamEvent::ThinkingSignature(sig) => thinking_signature = Some(sig),
163 LlmStreamEvent::ReasonItem {
164 encrypted_content, ..
165 } => {
166 if let Some(sig) = encrypted_content {
167 thinking_signature = Some(sig);
168 }
169 }
170 LlmStreamEvent::ToolCalls(calls) => tool_calls = calls,
171 LlmStreamEvent::Done(meta) => metadata = *meta,
172 LlmStreamEvent::Error(err) => return Err(crate::error::AgentLoopError::llm(err)),
173 }
174 }
175
176 Ok(LlmResponse {
177 text,
178 thinking: if thinking.is_empty() {
179 None
180 } else {
181 Some(thinking)
182 },
183 thinking_signature,
184 tool_calls: if tool_calls.is_empty() {
185 None
186 } else {
187 Some(tool_calls)
188 },
189 metadata,
190 })
191 }
192
193 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
201 Ok(None)
203 }
204
205 fn supports_compact(&self) -> bool {
214 false
216 }
217
218 async fn compact(&self, _request: CompactRequest) -> Result<Option<CompactResponse>> {
238 Ok(None)
240 }
241}
242
243#[async_trait]
245impl LlmDriver for Box<dyn LlmDriver> {
246 async fn chat_completion_stream(
247 &self,
248 messages: Vec<LlmMessage>,
249 config: &LlmCallConfig,
250 ) -> Result<LlmResponseStream> {
251 (**self).chat_completion_stream(messages, config).await
252 }
253
254 async fn chat_completion(
255 &self,
256 messages: Vec<LlmMessage>,
257 config: &LlmCallConfig,
258 ) -> Result<LlmResponse> {
259 (**self).chat_completion(messages, config).await
260 }
261
262 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
263 (**self).list_models().await
264 }
265
266 fn supports_compact(&self) -> bool {
267 (**self).supports_compact()
268 }
269
270 async fn compact(&self, request: CompactRequest) -> Result<Option<CompactResponse>> {
271 (**self).compact(request).await
272 }
273}
274
275#[derive(Debug, Clone)]
281pub struct LlmMessage {
282 pub role: LlmMessageRole,
283 pub content: LlmMessageContent,
284 pub tool_calls: Option<Vec<ToolCall>>,
285 pub tool_call_id: Option<String>,
286 pub phase: Option<crate::message::ExecutionPhase>,
291 pub thinking: Option<String>,
294 pub thinking_signature: Option<String>,
297}
298
299impl LlmMessage {
300 pub fn text(role: LlmMessageRole, content: impl Into<String>) -> Self {
302 Self {
303 role,
304 content: LlmMessageContent::Text(content.into()),
305 tool_calls: None,
306 tool_call_id: None,
307 phase: None,
308 thinking: None,
309 thinking_signature: None,
310 }
311 }
312
313 pub fn parts(role: LlmMessageRole, parts: Vec<LlmContentPart>) -> Self {
315 Self {
316 role,
317 content: LlmMessageContent::Parts(parts),
318 tool_calls: None,
319 tool_call_id: None,
320 phase: None,
321 thinking: None,
322 thinking_signature: None,
323 }
324 }
325
326 pub fn content_as_text(&self) -> String {
328 self.content.to_text()
329 }
330
331 pub fn prepend_text_prefix(&mut self, prefix: &str) {
336 match &mut self.content {
337 LlmMessageContent::Text(text) => {
338 *text = format!("{}{}", prefix, text);
339 }
340 LlmMessageContent::Parts(parts) => {
341 for part in parts.iter_mut() {
342 if let LlmContentPart::Text { text } = part {
343 *text = format!("{}{}", prefix, text);
344 return;
345 }
346 }
347 parts.insert(
349 0,
350 LlmContentPart::Text {
351 text: prefix.to_string(),
352 },
353 );
354 }
355 }
356 }
357}
358
359#[derive(Debug, Clone)]
361pub enum LlmMessageContent {
362 Text(String),
364 Parts(Vec<LlmContentPart>),
366}
367
368impl LlmMessageContent {
369 pub fn to_text(&self) -> String {
371 match self {
372 LlmMessageContent::Text(s) => s.clone(),
373 LlmMessageContent::Parts(parts) => parts
374 .iter()
375 .filter_map(|p| match p {
376 LlmContentPart::Text { text } => Some(text.clone()),
377 _ => None,
378 })
379 .collect::<Vec<_>>()
380 .join(""),
381 }
382 }
383
384 pub fn is_text(&self) -> bool {
386 matches!(self, LlmMessageContent::Text(_))
387 }
388
389 pub fn is_parts(&self) -> bool {
391 matches!(self, LlmMessageContent::Parts(_))
392 }
393}
394
395impl From<String> for LlmMessageContent {
396 fn from(s: String) -> Self {
397 LlmMessageContent::Text(s)
398 }
399}
400
401impl From<&str> for LlmMessageContent {
402 fn from(s: &str) -> Self {
403 LlmMessageContent::Text(s.to_string())
404 }
405}
406
407#[derive(Debug, Clone)]
409pub enum LlmContentPart {
410 Text { text: String },
412 Image { url: String },
414 Audio { url: String },
416}
417
418impl LlmContentPart {
419 pub fn text(text: impl Into<String>) -> Self {
421 LlmContentPart::Text { text: text.into() }
422 }
423
424 pub fn image(url: impl Into<String>) -> Self {
426 LlmContentPart::Image { url: url.into() }
427 }
428
429 pub fn audio(url: impl Into<String>) -> Self {
431 LlmContentPart::Audio { url: url.into() }
432 }
433}
434
435#[derive(Debug, Clone, PartialEq, Eq)]
437pub enum LlmMessageRole {
438 System,
439 User,
440 Assistant,
441 Tool,
442}
443
444#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
454pub struct ToolSearchConfig {
455 pub enabled: bool,
457 pub threshold: usize,
460}
461
462#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
464#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
465#[serde(rename_all = "snake_case")]
466pub enum PromptCacheStrategy {
467 #[default]
469 Auto,
470}
471
472#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
477#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
478pub struct PromptCacheConfig {
479 pub enabled: bool,
481 #[serde(default)]
483 pub strategy: PromptCacheStrategy,
484 #[serde(default, skip_serializing_if = "Option::is_none")]
491 pub gemini_cached_content: Option<String>,
492}
493
494#[derive(Debug, Clone)]
496pub struct LlmCallConfig {
497 pub model: String,
498 pub temperature: Option<f32>,
499 pub max_tokens: Option<u32>,
500 pub tools: Vec<ToolDefinition>,
501 pub reasoning_effort: Option<String>,
503 pub metadata: HashMap<String, String>,
507 pub previous_response_id: Option<String>,
510 pub tool_search: Option<ToolSearchConfig>,
512 pub prompt_cache: Option<PromptCacheConfig>,
514}
515
516impl From<&RuntimeAgent> for LlmCallConfig {
517 fn from(runtime_agent: &RuntimeAgent) -> Self {
518 Self {
519 model: runtime_agent.model.clone(),
520 temperature: runtime_agent.temperature,
521 max_tokens: runtime_agent.max_tokens,
522 tools: runtime_agent.tools.clone(),
523 reasoning_effort: None, metadata: HashMap::new(), previous_response_id: None,
526 tool_search: runtime_agent.tool_search.clone(),
527 prompt_cache: runtime_agent.prompt_cache.clone(),
528 }
529 }
530}
531
532#[derive(Debug, Clone)]
534pub struct LlmResponse {
535 pub text: String,
536 pub thinking: Option<String>,
538 pub thinking_signature: Option<String>,
540 pub tool_calls: Option<Vec<ToolCall>>,
541 pub metadata: LlmCompletionMetadata,
542}
543
544pub struct LlmCallConfigBuilder {
563 config: LlmCallConfig,
564}
565
566impl LlmCallConfigBuilder {
567 pub fn from(runtime_agent: &RuntimeAgent) -> Self {
569 Self {
570 config: LlmCallConfig::from(runtime_agent),
571 }
572 }
573
574 pub fn reasoning_effort(mut self, effort: impl Into<String>) -> Self {
576 self.config.reasoning_effort = Some(effort.into());
577 self
578 }
579
580 pub fn model(mut self, model: impl Into<String>) -> Self {
582 self.config.model = model.into();
583 self
584 }
585
586 pub fn temperature(mut self, temp: f32) -> Self {
588 self.config.temperature = Some(temp);
589 self
590 }
591
592 pub fn max_tokens(mut self, tokens: u32) -> Self {
594 self.config.max_tokens = Some(tokens);
595 self
596 }
597
598 pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
600 self.config.tools = tools;
601 self
602 }
603
604 pub fn metadata(mut self, metadata: HashMap<String, String>) -> Self {
609 self.config.metadata = metadata;
610 self
611 }
612
613 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
615 self.config.metadata.insert(key.into(), value.into());
616 self
617 }
618
619 pub fn previous_response_id(mut self, id: Option<String>) -> Self {
621 self.config.previous_response_id = id;
622 self
623 }
624
625 pub fn tool_search(mut self, config: ToolSearchConfig) -> Self {
627 self.config.tool_search = Some(config);
628 self
629 }
630
631 pub fn prompt_cache(mut self, config: PromptCacheConfig) -> Self {
633 self.config.prompt_cache = Some(config);
634 self
635 }
636
637 pub fn build(self) -> LlmCallConfig {
639 self.config
640 }
641}
642
643impl From<&crate::message::Message> for LlmMessage {
648 fn from(msg: &crate::message::Message) -> Self {
654 let role = match msg.role {
655 crate::message::MessageRole::System => LlmMessageRole::System,
656 crate::message::MessageRole::User => LlmMessageRole::User,
657 crate::message::MessageRole::Agent => LlmMessageRole::Assistant,
658 crate::message::MessageRole::ToolResult => LlmMessageRole::Tool,
659 };
660
661 let tool_calls: Vec<ToolCall> = msg
663 .tool_calls()
664 .into_iter()
665 .map(|tc| ToolCall {
666 id: tc.id.clone(),
667 name: tc.name.clone(),
668 arguments: tc.arguments.clone(),
669 })
670 .collect();
671
672 LlmMessage {
673 role,
674 content: LlmMessageContent::Text(msg.content_to_llm_string()),
675 tool_calls: if tool_calls.is_empty() {
676 None
677 } else {
678 Some(tool_calls)
679 },
680 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
681 phase: msg.phase,
682 thinking: msg.thinking.clone(),
683 thinking_signature: msg.thinking_signature.clone(),
684 }
685 }
686}
687
688use crate::traits::ResolvedImage;
693use uuid::Uuid;
694
695impl LlmMessage {
696 pub fn from_message_with_images(
716 msg: &crate::message::Message,
717 resolved_images: &HashMap<Uuid, ResolvedImage>,
718 ) -> Self {
719 use crate::message::{ContentPart, MessageRole};
720
721 let role = match msg.role {
722 MessageRole::System => LlmMessageRole::System,
723 MessageRole::User => LlmMessageRole::User,
724 MessageRole::Agent => LlmMessageRole::Assistant,
725 MessageRole::ToolResult => LlmMessageRole::Tool,
726 };
727
728 let mut parts: Vec<LlmContentPart> = Vec::new();
730 let mut tool_calls: Vec<ToolCall> = Vec::new();
731
732 for part in &msg.content {
733 match part {
734 ContentPart::Text(t) => {
735 parts.push(LlmContentPart::Text {
736 text: t.text.clone(),
737 });
738 }
739 ContentPart::Image(img) => {
740 if let Some(url) = &img.url {
742 parts.push(LlmContentPart::Image { url: url.clone() });
743 } else if let (Some(base64), Some(media_type)) = (&img.base64, &img.media_type)
744 {
745 let data_url = format!("data:{};base64,{}", media_type, base64);
746 parts.push(LlmContentPart::Image { url: data_url });
747 }
748 }
749 ContentPart::ImageFile(img_file) => {
750 if let Some(resolved) = resolved_images.get(&img_file.image_id.uuid()) {
752 parts.push(LlmContentPart::Image {
753 url: resolved.to_data_url(),
754 });
755 } else {
756 parts.push(LlmContentPart::Text {
758 text: format!("[Image not found: {}]", img_file.image_id),
759 });
760 }
761 }
762 ContentPart::ToolCall(tc) => {
763 tool_calls.push(ToolCall {
765 id: tc.id.clone(),
766 name: tc.name.clone(),
767 arguments: tc.arguments.clone(),
768 });
769 }
770 ContentPart::ToolResult(tr) => {
771 let text = if let Some(err) = &tr.error {
773 format!("Tool error: {}", err)
774 } else if let Some(res) = &tr.result {
775 serde_json::to_string(res).unwrap_or_else(|_| "{}".to_string())
776 } else {
777 "{}".to_string()
778 };
779 let text = truncate_tool_result(text);
783 parts.push(LlmContentPart::Text { text });
784 }
785 }
786 }
787
788 let content = if parts.len() == 1 && matches!(&parts[0], LlmContentPart::Text { .. }) {
790 if let LlmContentPart::Text { text } = &parts[0] {
792 LlmMessageContent::Text(text.clone())
793 } else {
794 LlmMessageContent::Parts(parts)
795 }
796 } else if parts.is_empty() {
797 LlmMessageContent::Text(String::new())
799 } else {
800 LlmMessageContent::Parts(parts)
802 };
803
804 LlmMessage {
805 role,
806 content,
807 tool_calls: if tool_calls.is_empty() {
808 None
809 } else {
810 Some(tool_calls)
811 },
812 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
813 phase: msg.phase,
814 thinking: msg.thinking.clone(),
815 thinking_signature: msg.thinking_signature.clone(),
816 }
817 }
818
819 pub fn message_has_image_files(msg: &crate::message::Message) -> bool {
821 msg.content.iter().any(|p| p.is_image_file())
822 }
823
824 pub fn extract_image_file_ids(msg: &crate::message::Message) -> Vec<Uuid> {
826 msg.content
827 .iter()
828 .filter_map(|p| match p {
829 crate::message::ContentPart::ImageFile(f) => Some(f.image_id.uuid()),
830 _ => None,
831 })
832 .collect()
833 }
834}
835
836#[derive(Debug, Clone, PartialEq, Eq, Hash)]
842pub enum ProviderType {
843 OpenAI,
846 AzureOpenAI,
848 OpenAICompletions,
851 Anthropic,
852 Gemini,
854 LlmSim,
856}
857
858impl std::str::FromStr for ProviderType {
859 type Err = String;
860
861 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
862 match s.to_lowercase().as_str() {
863 "openai" => Ok(ProviderType::OpenAI),
864 "azure_openai" => Ok(ProviderType::AzureOpenAI),
865 "openai_completions" => Ok(ProviderType::OpenAICompletions),
866 "anthropic" => Ok(ProviderType::Anthropic),
867 "gemini" => Ok(ProviderType::Gemini),
868 "llmsim" => Ok(ProviderType::LlmSim),
869 _ => Err(format!("Unknown provider type: {}", s)),
870 }
871 }
872}
873
874impl std::fmt::Display for ProviderType {
875 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
876 match self {
877 ProviderType::OpenAI => write!(f, "openai"),
878 ProviderType::AzureOpenAI => write!(f, "azure_openai"),
879 ProviderType::OpenAICompletions => write!(f, "openai_completions"),
880 ProviderType::Anthropic => write!(f, "anthropic"),
881 ProviderType::Gemini => write!(f, "gemini"),
882 ProviderType::LlmSim => write!(f, "llmsim"),
883 }
884 }
885}
886
887#[derive(Debug, Clone)]
889pub struct ProviderConfig {
890 pub provider_type: ProviderType,
892 pub api_key: Option<String>,
894 pub base_url: Option<String>,
896}
897
898impl ProviderConfig {
899 pub fn new(provider_type: ProviderType) -> Self {
901 Self {
902 provider_type,
903 api_key: None,
904 base_url: None,
905 }
906 }
907
908 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
910 self.api_key = Some(api_key.into());
911 self
912 }
913
914 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
916 self.base_url = Some(base_url.into());
917 self
918 }
919}
920
921pub type BoxedLlmDriver = Box<dyn LlmDriver>;
923
924pub type DriverFactory = Arc<dyn Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync>;
932
933#[derive(Clone, Default)]
953pub struct DriverRegistry {
954 factories: HashMap<ProviderType, DriverFactory>,
955}
956
957impl DriverRegistry {
958 pub fn new() -> Self {
960 Self {
961 factories: HashMap::new(),
962 }
963 }
964
965 pub fn register<F>(&mut self, provider_type: ProviderType, factory: F)
967 where
968 F: Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync + 'static,
969 {
970 self.factories.insert(provider_type, Arc::new(factory));
971 }
972
973 pub fn create_driver(&self, config: &ProviderConfig) -> Result<BoxedLlmDriver> {
981 let api_key = if config.provider_type == ProviderType::LlmSim {
983 config.api_key.as_deref().unwrap_or("")
985 } else {
986 config.api_key.as_ref().ok_or_else(|| {
987 AgentLoopError::llm(
988 "API key is required. Configure the API key in provider settings.",
989 )
990 })?
991 };
992
993 let factory = self.factories.get(&config.provider_type).ok_or_else(|| {
995 AgentLoopError::driver_not_registered(config.provider_type.to_string())
996 })?;
997
998 Ok(factory(api_key, config.base_url.as_deref()))
1000 }
1001
1002 pub fn has_driver(&self, provider_type: &ProviderType) -> bool {
1004 self.factories.contains_key(provider_type)
1005 }
1006
1007 pub fn registered_providers(&self) -> Vec<ProviderType> {
1009 self.factories.keys().cloned().collect()
1010 }
1011}
1012
1013const MAX_TOOL_RESULT_BYTES: usize = 64 * 1024;
1018
1019const TRUNCATION_SUFFIX: &str =
1020 "\n\n[Output truncated — exceeded 64 KiB limit. Try quiet flags, pipes, or redirect to file.]";
1021
1022fn truncate_tool_result(text: String) -> String {
1023 if text.len() <= MAX_TOOL_RESULT_BYTES {
1024 return text;
1025 }
1026 let content_budget = MAX_TOOL_RESULT_BYTES.saturating_sub(TRUNCATION_SUFFIX.len());
1027 let mut end = content_budget;
1028 while end > 0 && !text.is_char_boundary(end) {
1029 end -= 1;
1030 }
1031 let mut truncated = text[..end].to_string();
1032 truncated.push_str(TRUNCATION_SUFFIX);
1033 truncated
1034}
1035
1036#[cfg(test)]
1041mod tests {
1042 use super::*;
1043
1044 #[test]
1045 fn test_llm_call_config_builder_from_runtime_agent() {
1046 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1047 let llm_config = LlmCallConfigBuilder::from(&runtime_agent).build();
1048
1049 assert_eq!(llm_config.model, "gpt-4o");
1050 assert!(llm_config.reasoning_effort.is_none());
1051 assert!(llm_config.temperature.is_none());
1052 assert!(llm_config.max_tokens.is_none());
1053 assert!(llm_config.tools.is_empty());
1054 assert!(llm_config.metadata.is_empty());
1055 }
1056
1057 #[test]
1058 fn test_llm_call_config_builder_with_metadata() {
1059 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1060 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1061 .with_metadata("session_id", "session_abc123")
1062 .with_metadata("agent_id", "agent_xyz789")
1063 .build();
1064
1065 assert_eq!(
1066 llm_config.metadata.get("session_id"),
1067 Some(&"session_abc123".to_string())
1068 );
1069 assert_eq!(
1070 llm_config.metadata.get("agent_id"),
1071 Some(&"agent_xyz789".to_string())
1072 );
1073 }
1074
1075 #[test]
1076 fn test_llm_call_config_builder_with_metadata_hashmap() {
1077 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1078 let mut metadata = HashMap::new();
1079 metadata.insert("key1".to_string(), "value1".to_string());
1080 metadata.insert("key2".to_string(), "value2".to_string());
1081
1082 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1083 .metadata(metadata)
1084 .build();
1085
1086 assert_eq!(llm_config.metadata.get("key1"), Some(&"value1".to_string()));
1087 assert_eq!(llm_config.metadata.get("key2"), Some(&"value2".to_string()));
1088 }
1089
1090 #[test]
1091 fn test_llm_call_config_builder_with_reasoning_effort() {
1092 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1093 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1094 .reasoning_effort("high")
1095 .build();
1096
1097 assert_eq!(llm_config.reasoning_effort, Some("high".to_string()));
1098 }
1099
1100 #[test]
1101 fn test_llm_call_config_builder_with_all_options() {
1102 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1103 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1104 .model("claude-3-opus")
1105 .reasoning_effort("medium")
1106 .temperature(0.7)
1107 .max_tokens(1000)
1108 .build();
1109
1110 assert_eq!(llm_config.model, "claude-3-opus");
1111 assert_eq!(llm_config.reasoning_effort, Some("medium".to_string()));
1112 assert_eq!(llm_config.temperature, Some(0.7));
1113 assert_eq!(llm_config.max_tokens, Some(1000));
1114 }
1115
1116 #[test]
1117 fn test_provider_type_parsing() {
1118 assert_eq!(
1119 "openai".parse::<ProviderType>().unwrap(),
1120 ProviderType::OpenAI
1121 );
1122 assert_eq!(
1123 "openai_completions".parse::<ProviderType>().unwrap(),
1124 ProviderType::OpenAICompletions
1125 );
1126 assert_eq!(
1127 "azure_openai".parse::<ProviderType>().unwrap(),
1128 ProviderType::AzureOpenAI
1129 );
1130 assert_eq!(
1131 "anthropic".parse::<ProviderType>().unwrap(),
1132 ProviderType::Anthropic
1133 );
1134 assert_eq!(
1135 "gemini".parse::<ProviderType>().unwrap(),
1136 ProviderType::Gemini
1137 );
1138 assert!("ollama".parse::<ProviderType>().is_err());
1140 assert!("custom".parse::<ProviderType>().is_err());
1141 }
1142
1143 #[test]
1144 fn test_provider_type_display() {
1145 assert_eq!(ProviderType::OpenAI.to_string(), "openai");
1146 assert_eq!(ProviderType::AzureOpenAI.to_string(), "azure_openai");
1147 assert_eq!(
1148 ProviderType::OpenAICompletions.to_string(),
1149 "openai_completions"
1150 );
1151 assert_eq!(ProviderType::Anthropic.to_string(), "anthropic");
1152 assert_eq!(ProviderType::Gemini.to_string(), "gemini");
1153 }
1154
1155 #[test]
1156 fn test_provider_config_builder() {
1157 let config = ProviderConfig::new(ProviderType::Anthropic)
1158 .with_api_key("test-key")
1159 .with_base_url("https://custom.api.com");
1160
1161 assert_eq!(config.provider_type, ProviderType::Anthropic);
1162 assert_eq!(config.api_key, Some("test-key".to_string()));
1163 assert_eq!(config.base_url, Some("https://custom.api.com".to_string()));
1164 }
1165
1166 #[test]
1167 fn test_driver_registry_requires_api_key() {
1168 let mut registry = DriverRegistry::new();
1170 registry.register(ProviderType::OpenAI, |_api_key, _base_url| {
1171 struct MockDriver;
1173 #[async_trait]
1174 impl LlmDriver for MockDriver {
1175 async fn chat_completion_stream(
1176 &self,
1177 _messages: Vec<LlmMessage>,
1178 _config: &LlmCallConfig,
1179 ) -> Result<LlmResponseStream> {
1180 unimplemented!()
1181 }
1182 }
1183 Box::new(MockDriver)
1184 });
1185
1186 let config = ProviderConfig::new(ProviderType::OpenAI);
1188 let result = registry.create_driver(&config);
1189 assert!(result.is_err());
1190
1191 let config_with_key = ProviderConfig::new(ProviderType::OpenAI).with_api_key("test-key");
1193 let result = registry.create_driver(&config_with_key);
1194 assert!(result.is_ok());
1195 }
1196
1197 #[test]
1198 fn test_driver_registry_returns_error_for_unregistered_provider() {
1199 let registry = DriverRegistry::new();
1200 let config = ProviderConfig::new(ProviderType::Anthropic).with_api_key("test-key");
1201
1202 let result = registry.create_driver(&config);
1203
1204 if let Err(AgentLoopError::DriverNotRegistered(provider)) = result {
1206 assert_eq!(provider, "anthropic");
1207 } else {
1208 panic!("Expected DriverNotRegistered error");
1209 }
1210 }
1211
1212 #[test]
1213 fn test_driver_registry_registration() {
1214 let mut registry = DriverRegistry::new();
1215
1216 assert!(!registry.has_driver(&ProviderType::OpenAI));
1217 assert!(!registry.has_driver(&ProviderType::Anthropic));
1218
1219 registry.register(ProviderType::OpenAI, |_, _| {
1220 struct MockDriver;
1221 #[async_trait]
1222 impl LlmDriver for MockDriver {
1223 async fn chat_completion_stream(
1224 &self,
1225 _messages: Vec<LlmMessage>,
1226 _config: &LlmCallConfig,
1227 ) -> Result<LlmResponseStream> {
1228 unimplemented!()
1229 }
1230 }
1231 Box::new(MockDriver)
1232 });
1233
1234 assert!(registry.has_driver(&ProviderType::OpenAI));
1235 assert!(!registry.has_driver(&ProviderType::Anthropic));
1236 }
1237
1238 use crate::{ContentPart, ImageFileContentPart, Message, MessageRole, TextContentPart};
1243
1244 #[test]
1245 fn test_message_has_image_files_with_image_file() {
1246 let message = Message {
1247 id: uuid::Uuid::new_v4().into(),
1248 role: MessageRole::User,
1249 content: vec![
1250 ContentPart::Text(TextContentPart {
1251 text: "Look at this image".to_string(),
1252 }),
1253 ContentPart::ImageFile(ImageFileContentPart {
1254 image_id: uuid::Uuid::new_v4().into(),
1255 filename: Some("test.png".to_string()),
1256 }),
1257 ],
1258 phase: None,
1259 thinking: None,
1260 thinking_signature: None,
1261 controls: None,
1262 metadata: None,
1263 external_actor: None,
1264 created_at: chrono::Utc::now(),
1265 };
1266
1267 assert!(LlmMessage::message_has_image_files(&message));
1268 }
1269
1270 #[test]
1271 fn test_message_has_image_files_without_image_file() {
1272 let message = Message {
1273 id: uuid::Uuid::new_v4().into(),
1274 role: MessageRole::User,
1275 content: vec![ContentPart::Text(TextContentPart {
1276 text: "Just text".to_string(),
1277 })],
1278 phase: None,
1279 thinking: None,
1280 thinking_signature: None,
1281 controls: None,
1282 metadata: None,
1283 external_actor: None,
1284 created_at: chrono::Utc::now(),
1285 };
1286
1287 assert!(!LlmMessage::message_has_image_files(&message));
1288 }
1289
1290 #[test]
1291 fn test_extract_image_file_ids() {
1292 let id1 = uuid::Uuid::new_v4();
1293 let id2 = uuid::Uuid::new_v4();
1294
1295 let message = Message {
1296 id: uuid::Uuid::new_v4().into(),
1297 role: MessageRole::User,
1298 content: vec![
1299 ContentPart::Text(TextContentPart {
1300 text: "Look at these images".to_string(),
1301 }),
1302 ContentPart::ImageFile(ImageFileContentPart {
1303 image_id: id1.into(),
1304 filename: Some("test1.png".to_string()),
1305 }),
1306 ContentPart::ImageFile(ImageFileContentPart {
1307 image_id: id2.into(),
1308 filename: Some("test2.png".to_string()),
1309 }),
1310 ],
1311 phase: None,
1312 thinking: None,
1313 thinking_signature: None,
1314 controls: None,
1315 metadata: None,
1316 external_actor: None,
1317 created_at: chrono::Utc::now(),
1318 };
1319
1320 let ids = LlmMessage::extract_image_file_ids(&message);
1321 assert_eq!(ids.len(), 2);
1322 assert!(ids.contains(&id1));
1323 assert!(ids.contains(&id2));
1324 }
1325
1326 #[test]
1327 fn test_from_message_with_images_text_only() {
1328 let message = Message {
1329 id: uuid::Uuid::new_v4().into(),
1330 role: MessageRole::User,
1331 content: vec![ContentPart::Text(TextContentPart {
1332 text: "Hello".to_string(),
1333 })],
1334 phase: None,
1335 thinking: None,
1336 thinking_signature: None,
1337 controls: None,
1338 metadata: None,
1339 external_actor: None,
1340 created_at: chrono::Utc::now(),
1341 };
1342
1343 let resolved = std::collections::HashMap::new();
1344 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1345
1346 assert_eq!(llm_message.role, LlmMessageRole::User);
1347 match llm_message.content {
1348 LlmMessageContent::Text(text) => assert_eq!(text, "Hello"),
1349 _ => panic!("Expected text content"),
1350 }
1351 }
1352
1353 #[test]
1354 fn test_from_message_with_images_resolved_image() {
1355 let image_id = uuid::Uuid::new_v4();
1356 let message = Message {
1357 id: uuid::Uuid::new_v4().into(),
1358 role: MessageRole::User,
1359 content: vec![
1360 ContentPart::Text(TextContentPart {
1361 text: "Look at this".to_string(),
1362 }),
1363 ContentPart::ImageFile(ImageFileContentPart {
1364 image_id: image_id.into(),
1365 filename: Some("test.png".to_string()),
1366 }),
1367 ],
1368 phase: None,
1369 thinking: None,
1370 thinking_signature: None,
1371 controls: None,
1372 metadata: None,
1373 external_actor: None,
1374 created_at: chrono::Utc::now(),
1375 };
1376
1377 let mut resolved = std::collections::HashMap::new();
1378 resolved.insert(
1379 image_id,
1380 crate::ResolvedImage::new("base64data", "image/png"),
1381 );
1382
1383 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1384
1385 match &llm_message.content {
1386 LlmMessageContent::Parts(parts) => {
1387 assert_eq!(parts.len(), 2);
1388 assert!(matches!(&parts[0], LlmContentPart::Text { .. }));
1390 if let LlmContentPart::Image { url } = &parts[1] {
1392 assert!(url.starts_with("data:image/png;base64,"));
1393 } else {
1394 panic!("Expected image content part");
1395 }
1396 }
1397 _ => panic!("Expected parts content"),
1398 }
1399 }
1400
1401 #[test]
1402 fn test_from_message_with_images_unresolved_image() {
1403 let image_id = uuid::Uuid::new_v4();
1404 let message = Message {
1405 id: uuid::Uuid::new_v4().into(),
1406 role: MessageRole::User,
1407 content: vec![ContentPart::ImageFile(ImageFileContentPart {
1408 image_id: image_id.into(),
1409 filename: Some("missing.png".to_string()),
1410 })],
1411 phase: None,
1412 thinking: None,
1413 thinking_signature: None,
1414 controls: None,
1415 metadata: None,
1416 external_actor: None,
1417 created_at: chrono::Utc::now(),
1418 };
1419
1420 let resolved = std::collections::HashMap::new();
1422 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1423
1424 match &llm_message.content {
1427 LlmMessageContent::Text(text) => {
1428 assert!(text.contains("Image not found"));
1429 }
1430 LlmMessageContent::Parts(parts) => {
1431 assert_eq!(parts.len(), 1);
1432 if let LlmContentPart::Text { text } = &parts[0] {
1433 assert!(text.contains("Image not found"));
1434 } else {
1435 panic!("Expected text placeholder for missing image");
1436 }
1437 }
1438 }
1439 }
1440
1441 #[test]
1442 fn test_prepend_text_prefix_simple_text() {
1443 let mut msg = LlmMessage::text(LlmMessageRole::User, "Hello bot");
1444 msg.prepend_text_prefix("[Alice] ");
1445 assert_eq!(msg.content_as_text(), "[Alice] Hello bot");
1446 }
1447
1448 #[test]
1449 fn test_prepend_text_prefix_parts() {
1450 let mut msg = LlmMessage::parts(
1451 LlmMessageRole::User,
1452 vec![
1453 LlmContentPart::Text {
1454 text: "Hello".to_string(),
1455 },
1456 LlmContentPart::Image {
1457 url: "data:image/png;base64,abc".to_string(),
1458 },
1459 ],
1460 );
1461 msg.prepend_text_prefix("[Bob] ");
1462 match &msg.content {
1463 LlmMessageContent::Parts(parts) => {
1464 if let LlmContentPart::Text { text } = &parts[0] {
1465 assert_eq!(text, "[Bob] Hello");
1466 } else {
1467 panic!("Expected text part");
1468 }
1469 }
1470 _ => panic!("Expected parts content"),
1471 }
1472 }
1473
1474 #[test]
1475 fn test_prepend_text_prefix_parts_no_text() {
1476 let mut msg = LlmMessage::parts(
1477 LlmMessageRole::User,
1478 vec![LlmContentPart::Image {
1479 url: "data:image/png;base64,abc".to_string(),
1480 }],
1481 );
1482 msg.prepend_text_prefix("[Eve] ");
1483 match &msg.content {
1484 LlmMessageContent::Parts(parts) => {
1485 assert_eq!(parts.len(), 2);
1486 if let LlmContentPart::Text { text } = &parts[0] {
1487 assert_eq!(text, "[Eve] ");
1488 } else {
1489 panic!("Expected prepended text part");
1490 }
1491 }
1492 _ => panic!("Expected parts content"),
1493 }
1494 }
1495}