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 OpenRouter,
852 AzureOpenAI,
854 OpenAICompletions,
857 Anthropic,
858 Gemini,
860 LlmSim,
862 Bedrock,
864}
865
866impl std::str::FromStr for ProviderType {
867 type Err = String;
868
869 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
870 match s.to_lowercase().as_str() {
871 "openai" => Ok(ProviderType::OpenAI),
872 "openrouter" => Ok(ProviderType::OpenRouter),
873 "azure_openai" => Ok(ProviderType::AzureOpenAI),
874 "openai_completions" => Ok(ProviderType::OpenAICompletions),
875 "anthropic" => Ok(ProviderType::Anthropic),
876 "gemini" => Ok(ProviderType::Gemini),
877 "llmsim" => Ok(ProviderType::LlmSim),
878 "bedrock" => Ok(ProviderType::Bedrock),
879 _ => Err(format!("Unknown provider type: {}", s)),
880 }
881 }
882}
883
884impl std::fmt::Display for ProviderType {
885 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
886 match self {
887 ProviderType::OpenAI => write!(f, "openai"),
888 ProviderType::OpenRouter => write!(f, "openrouter"),
889 ProviderType::AzureOpenAI => write!(f, "azure_openai"),
890 ProviderType::OpenAICompletions => write!(f, "openai_completions"),
891 ProviderType::Anthropic => write!(f, "anthropic"),
892 ProviderType::Gemini => write!(f, "gemini"),
893 ProviderType::LlmSim => write!(f, "llmsim"),
894 ProviderType::Bedrock => write!(f, "bedrock"),
895 }
896 }
897}
898
899#[derive(Debug, Clone)]
901pub struct ProviderConfig {
902 pub provider_type: ProviderType,
904 pub api_key: Option<String>,
906 pub base_url: Option<String>,
908}
909
910impl ProviderConfig {
911 pub fn new(provider_type: ProviderType) -> Self {
913 Self {
914 provider_type,
915 api_key: None,
916 base_url: None,
917 }
918 }
919
920 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
922 self.api_key = Some(api_key.into());
923 self
924 }
925
926 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
928 self.base_url = Some(base_url.into());
929 self
930 }
931}
932
933pub type BoxedLlmDriver = Box<dyn LlmDriver>;
935
936pub type DriverFactory = Arc<dyn Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync>;
944
945#[derive(Clone, Default)]
965pub struct DriverRegistry {
966 factories: HashMap<ProviderType, DriverFactory>,
967}
968
969impl DriverRegistry {
970 pub fn new() -> Self {
972 Self {
973 factories: HashMap::new(),
974 }
975 }
976
977 pub fn register<F>(&mut self, provider_type: ProviderType, factory: F)
979 where
980 F: Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync + 'static,
981 {
982 self.factories.insert(provider_type, Arc::new(factory));
983 }
984
985 pub fn create_driver(&self, config: &ProviderConfig) -> Result<BoxedLlmDriver> {
993 let api_key = if config.provider_type == ProviderType::LlmSim {
995 config.api_key.as_deref().unwrap_or("")
997 } else {
998 config.api_key.as_ref().ok_or_else(|| {
999 AgentLoopError::llm(
1000 "API key is required. Configure the API key in provider settings.",
1001 )
1002 })?
1003 };
1004
1005 let factory = self.factories.get(&config.provider_type).ok_or_else(|| {
1007 AgentLoopError::driver_not_registered(config.provider_type.to_string())
1008 })?;
1009
1010 Ok(factory(api_key, config.base_url.as_deref()))
1012 }
1013
1014 pub fn has_driver(&self, provider_type: &ProviderType) -> bool {
1016 self.factories.contains_key(provider_type)
1017 }
1018
1019 pub fn registered_providers(&self) -> Vec<ProviderType> {
1021 self.factories.keys().cloned().collect()
1022 }
1023}
1024
1025const MAX_TOOL_RESULT_BYTES: usize = 64 * 1024;
1030
1031const TRUNCATION_SUFFIX: &str =
1032 "\n\n[Output truncated — exceeded 64 KiB limit. Try quiet flags, pipes, or redirect to file.]";
1033
1034fn truncate_tool_result(text: String) -> String {
1035 if text.len() <= MAX_TOOL_RESULT_BYTES {
1036 return text;
1037 }
1038 let content_budget = MAX_TOOL_RESULT_BYTES.saturating_sub(TRUNCATION_SUFFIX.len());
1039 let mut end = content_budget;
1040 while end > 0 && !text.is_char_boundary(end) {
1041 end -= 1;
1042 }
1043 let mut truncated = text[..end].to_string();
1044 truncated.push_str(TRUNCATION_SUFFIX);
1045 truncated
1046}
1047
1048#[cfg(test)]
1053mod tests {
1054 use super::*;
1055
1056 #[test]
1057 fn test_llm_call_config_builder_from_runtime_agent() {
1058 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1059 let llm_config = LlmCallConfigBuilder::from(&runtime_agent).build();
1060
1061 assert_eq!(llm_config.model, "gpt-4o");
1062 assert!(llm_config.reasoning_effort.is_none());
1063 assert!(llm_config.temperature.is_none());
1064 assert!(llm_config.max_tokens.is_none());
1065 assert!(llm_config.tools.is_empty());
1066 assert!(llm_config.metadata.is_empty());
1067 }
1068
1069 #[test]
1070 fn test_llm_call_config_builder_with_metadata() {
1071 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1072 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1073 .with_metadata("session_id", "session_abc123")
1074 .with_metadata("agent_id", "agent_xyz789")
1075 .build();
1076
1077 assert_eq!(
1078 llm_config.metadata.get("session_id"),
1079 Some(&"session_abc123".to_string())
1080 );
1081 assert_eq!(
1082 llm_config.metadata.get("agent_id"),
1083 Some(&"agent_xyz789".to_string())
1084 );
1085 }
1086
1087 #[test]
1088 fn test_llm_call_config_builder_with_metadata_hashmap() {
1089 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1090 let mut metadata = HashMap::new();
1091 metadata.insert("key1".to_string(), "value1".to_string());
1092 metadata.insert("key2".to_string(), "value2".to_string());
1093
1094 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1095 .metadata(metadata)
1096 .build();
1097
1098 assert_eq!(llm_config.metadata.get("key1"), Some(&"value1".to_string()));
1099 assert_eq!(llm_config.metadata.get("key2"), Some(&"value2".to_string()));
1100 }
1101
1102 #[test]
1103 fn test_llm_call_config_builder_with_reasoning_effort() {
1104 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1105 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1106 .reasoning_effort("high")
1107 .build();
1108
1109 assert_eq!(llm_config.reasoning_effort, Some("high".to_string()));
1110 }
1111
1112 #[test]
1113 fn test_llm_call_config_builder_with_all_options() {
1114 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1115 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1116 .model("claude-3-opus")
1117 .reasoning_effort("medium")
1118 .temperature(0.7)
1119 .max_tokens(1000)
1120 .build();
1121
1122 assert_eq!(llm_config.model, "claude-3-opus");
1123 assert_eq!(llm_config.reasoning_effort, Some("medium".to_string()));
1124 assert_eq!(llm_config.temperature, Some(0.7));
1125 assert_eq!(llm_config.max_tokens, Some(1000));
1126 }
1127
1128 #[test]
1129 fn test_provider_type_parsing() {
1130 assert_eq!(
1131 "openai".parse::<ProviderType>().unwrap(),
1132 ProviderType::OpenAI
1133 );
1134 assert_eq!(
1135 "openrouter".parse::<ProviderType>().unwrap(),
1136 ProviderType::OpenRouter
1137 );
1138 assert_eq!(
1139 "openai_completions".parse::<ProviderType>().unwrap(),
1140 ProviderType::OpenAICompletions
1141 );
1142 assert_eq!(
1143 "azure_openai".parse::<ProviderType>().unwrap(),
1144 ProviderType::AzureOpenAI
1145 );
1146 assert_eq!(
1147 "anthropic".parse::<ProviderType>().unwrap(),
1148 ProviderType::Anthropic
1149 );
1150 assert_eq!(
1151 "gemini".parse::<ProviderType>().unwrap(),
1152 ProviderType::Gemini
1153 );
1154 assert!("ollama".parse::<ProviderType>().is_err());
1156 assert!("custom".parse::<ProviderType>().is_err());
1157 }
1158
1159 #[test]
1160 fn test_provider_type_display() {
1161 assert_eq!(ProviderType::OpenAI.to_string(), "openai");
1162 assert_eq!(ProviderType::OpenRouter.to_string(), "openrouter");
1163 assert_eq!(ProviderType::AzureOpenAI.to_string(), "azure_openai");
1164 assert_eq!(
1165 ProviderType::OpenAICompletions.to_string(),
1166 "openai_completions"
1167 );
1168 assert_eq!(ProviderType::Anthropic.to_string(), "anthropic");
1169 assert_eq!(ProviderType::Gemini.to_string(), "gemini");
1170 }
1171
1172 #[test]
1173 fn test_provider_config_builder() {
1174 let config = ProviderConfig::new(ProviderType::Anthropic)
1175 .with_api_key("test-key")
1176 .with_base_url("https://custom.api.com");
1177
1178 assert_eq!(config.provider_type, ProviderType::Anthropic);
1179 assert_eq!(config.api_key, Some("test-key".to_string()));
1180 assert_eq!(config.base_url, Some("https://custom.api.com".to_string()));
1181 }
1182
1183 #[test]
1184 fn test_driver_registry_requires_api_key() {
1185 let mut registry = DriverRegistry::new();
1187 registry.register(ProviderType::OpenAI, |_api_key, _base_url| {
1188 struct MockDriver;
1190 #[async_trait]
1191 impl LlmDriver for MockDriver {
1192 async fn chat_completion_stream(
1193 &self,
1194 _messages: Vec<LlmMessage>,
1195 _config: &LlmCallConfig,
1196 ) -> Result<LlmResponseStream> {
1197 unimplemented!()
1198 }
1199 }
1200 Box::new(MockDriver)
1201 });
1202
1203 let config = ProviderConfig::new(ProviderType::OpenAI);
1205 let result = registry.create_driver(&config);
1206 assert!(result.is_err());
1207
1208 let config_with_key = ProviderConfig::new(ProviderType::OpenAI).with_api_key("test-key");
1210 let result = registry.create_driver(&config_with_key);
1211 assert!(result.is_ok());
1212 }
1213
1214 #[test]
1215 fn test_driver_registry_returns_error_for_unregistered_provider() {
1216 let registry = DriverRegistry::new();
1217 let config = ProviderConfig::new(ProviderType::Anthropic).with_api_key("test-key");
1218
1219 let result = registry.create_driver(&config);
1220
1221 if let Err(AgentLoopError::DriverNotRegistered(provider)) = result {
1223 assert_eq!(provider, "anthropic");
1224 } else {
1225 panic!("Expected DriverNotRegistered error");
1226 }
1227 }
1228
1229 #[test]
1230 fn test_driver_registry_registration() {
1231 let mut registry = DriverRegistry::new();
1232
1233 assert!(!registry.has_driver(&ProviderType::OpenAI));
1234 assert!(!registry.has_driver(&ProviderType::Anthropic));
1235
1236 registry.register(ProviderType::OpenAI, |_, _| {
1237 struct MockDriver;
1238 #[async_trait]
1239 impl LlmDriver for MockDriver {
1240 async fn chat_completion_stream(
1241 &self,
1242 _messages: Vec<LlmMessage>,
1243 _config: &LlmCallConfig,
1244 ) -> Result<LlmResponseStream> {
1245 unimplemented!()
1246 }
1247 }
1248 Box::new(MockDriver)
1249 });
1250
1251 assert!(registry.has_driver(&ProviderType::OpenAI));
1252 assert!(!registry.has_driver(&ProviderType::Anthropic));
1253 }
1254
1255 use crate::{ContentPart, ImageFileContentPart, Message, MessageRole, TextContentPart};
1260
1261 #[test]
1262 fn test_message_has_image_files_with_image_file() {
1263 let message = Message {
1264 id: uuid::Uuid::new_v4().into(),
1265 role: MessageRole::User,
1266 content: vec![
1267 ContentPart::Text(TextContentPart {
1268 text: "Look at this image".to_string(),
1269 }),
1270 ContentPart::ImageFile(ImageFileContentPart {
1271 image_id: uuid::Uuid::new_v4().into(),
1272 filename: Some("test.png".to_string()),
1273 }),
1274 ],
1275 phase: None,
1276 thinking: None,
1277 thinking_signature: None,
1278 controls: None,
1279 metadata: None,
1280 external_actor: None,
1281 created_at: chrono::Utc::now(),
1282 };
1283
1284 assert!(LlmMessage::message_has_image_files(&message));
1285 }
1286
1287 #[test]
1288 fn test_message_has_image_files_without_image_file() {
1289 let message = Message {
1290 id: uuid::Uuid::new_v4().into(),
1291 role: MessageRole::User,
1292 content: vec![ContentPart::Text(TextContentPart {
1293 text: "Just text".to_string(),
1294 })],
1295 phase: None,
1296 thinking: None,
1297 thinking_signature: None,
1298 controls: None,
1299 metadata: None,
1300 external_actor: None,
1301 created_at: chrono::Utc::now(),
1302 };
1303
1304 assert!(!LlmMessage::message_has_image_files(&message));
1305 }
1306
1307 #[test]
1308 fn test_extract_image_file_ids() {
1309 let id1 = uuid::Uuid::new_v4();
1310 let id2 = uuid::Uuid::new_v4();
1311
1312 let message = Message {
1313 id: uuid::Uuid::new_v4().into(),
1314 role: MessageRole::User,
1315 content: vec![
1316 ContentPart::Text(TextContentPart {
1317 text: "Look at these images".to_string(),
1318 }),
1319 ContentPart::ImageFile(ImageFileContentPart {
1320 image_id: id1.into(),
1321 filename: Some("test1.png".to_string()),
1322 }),
1323 ContentPart::ImageFile(ImageFileContentPart {
1324 image_id: id2.into(),
1325 filename: Some("test2.png".to_string()),
1326 }),
1327 ],
1328 phase: None,
1329 thinking: None,
1330 thinking_signature: None,
1331 controls: None,
1332 metadata: None,
1333 external_actor: None,
1334 created_at: chrono::Utc::now(),
1335 };
1336
1337 let ids = LlmMessage::extract_image_file_ids(&message);
1338 assert_eq!(ids.len(), 2);
1339 assert!(ids.contains(&id1));
1340 assert!(ids.contains(&id2));
1341 }
1342
1343 #[test]
1344 fn test_from_message_with_images_text_only() {
1345 let message = Message {
1346 id: uuid::Uuid::new_v4().into(),
1347 role: MessageRole::User,
1348 content: vec![ContentPart::Text(TextContentPart {
1349 text: "Hello".to_string(),
1350 })],
1351 phase: None,
1352 thinking: None,
1353 thinking_signature: None,
1354 controls: None,
1355 metadata: None,
1356 external_actor: None,
1357 created_at: chrono::Utc::now(),
1358 };
1359
1360 let resolved = std::collections::HashMap::new();
1361 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1362
1363 assert_eq!(llm_message.role, LlmMessageRole::User);
1364 match llm_message.content {
1365 LlmMessageContent::Text(text) => assert_eq!(text, "Hello"),
1366 _ => panic!("Expected text content"),
1367 }
1368 }
1369
1370 #[test]
1371 fn test_from_message_with_images_resolved_image() {
1372 let image_id = uuid::Uuid::new_v4();
1373 let message = Message {
1374 id: uuid::Uuid::new_v4().into(),
1375 role: MessageRole::User,
1376 content: vec![
1377 ContentPart::Text(TextContentPart {
1378 text: "Look at this".to_string(),
1379 }),
1380 ContentPart::ImageFile(ImageFileContentPart {
1381 image_id: image_id.into(),
1382 filename: Some("test.png".to_string()),
1383 }),
1384 ],
1385 phase: None,
1386 thinking: None,
1387 thinking_signature: None,
1388 controls: None,
1389 metadata: None,
1390 external_actor: None,
1391 created_at: chrono::Utc::now(),
1392 };
1393
1394 let mut resolved = std::collections::HashMap::new();
1395 resolved.insert(
1396 image_id,
1397 crate::ResolvedImage::new("base64data", "image/png"),
1398 );
1399
1400 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1401
1402 match &llm_message.content {
1403 LlmMessageContent::Parts(parts) => {
1404 assert_eq!(parts.len(), 2);
1405 assert!(matches!(&parts[0], LlmContentPart::Text { .. }));
1407 if let LlmContentPart::Image { url } = &parts[1] {
1409 assert!(url.starts_with("data:image/png;base64,"));
1410 } else {
1411 panic!("Expected image content part");
1412 }
1413 }
1414 _ => panic!("Expected parts content"),
1415 }
1416 }
1417
1418 #[test]
1419 fn test_from_message_with_images_unresolved_image() {
1420 let image_id = uuid::Uuid::new_v4();
1421 let message = Message {
1422 id: uuid::Uuid::new_v4().into(),
1423 role: MessageRole::User,
1424 content: vec![ContentPart::ImageFile(ImageFileContentPart {
1425 image_id: image_id.into(),
1426 filename: Some("missing.png".to_string()),
1427 })],
1428 phase: None,
1429 thinking: None,
1430 thinking_signature: None,
1431 controls: None,
1432 metadata: None,
1433 external_actor: None,
1434 created_at: chrono::Utc::now(),
1435 };
1436
1437 let resolved = std::collections::HashMap::new();
1439 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1440
1441 match &llm_message.content {
1444 LlmMessageContent::Text(text) => {
1445 assert!(text.contains("Image not found"));
1446 }
1447 LlmMessageContent::Parts(parts) => {
1448 assert_eq!(parts.len(), 1);
1449 if let LlmContentPart::Text { text } = &parts[0] {
1450 assert!(text.contains("Image not found"));
1451 } else {
1452 panic!("Expected text placeholder for missing image");
1453 }
1454 }
1455 }
1456 }
1457
1458 #[test]
1459 fn test_prepend_text_prefix_simple_text() {
1460 let mut msg = LlmMessage::text(LlmMessageRole::User, "Hello bot");
1461 msg.prepend_text_prefix("[Alice] ");
1462 assert_eq!(msg.content_as_text(), "[Alice] Hello bot");
1463 }
1464
1465 #[test]
1466 fn test_prepend_text_prefix_parts() {
1467 let mut msg = LlmMessage::parts(
1468 LlmMessageRole::User,
1469 vec![
1470 LlmContentPart::Text {
1471 text: "Hello".to_string(),
1472 },
1473 LlmContentPart::Image {
1474 url: "data:image/png;base64,abc".to_string(),
1475 },
1476 ],
1477 );
1478 msg.prepend_text_prefix("[Bob] ");
1479 match &msg.content {
1480 LlmMessageContent::Parts(parts) => {
1481 if let LlmContentPart::Text { text } = &parts[0] {
1482 assert_eq!(text, "[Bob] Hello");
1483 } else {
1484 panic!("Expected text part");
1485 }
1486 }
1487 _ => panic!("Expected parts content"),
1488 }
1489 }
1490
1491 #[test]
1492 fn test_prepend_text_prefix_parts_no_text() {
1493 let mut msg = LlmMessage::parts(
1494 LlmMessageRole::User,
1495 vec![LlmContentPart::Image {
1496 url: "data:image/png;base64,abc".to_string(),
1497 }],
1498 );
1499 msg.prepend_text_prefix("[Eve] ");
1500 match &msg.content {
1501 LlmMessageContent::Parts(parts) => {
1502 assert_eq!(parts.len(), 2);
1503 if let LlmContentPart::Text { text } = &parts[0] {
1504 assert_eq!(text, "[Eve] ");
1505 } else {
1506 panic!("Expected prepended text part");
1507 }
1508 }
1509 _ => panic!("Expected parts content"),
1510 }
1511 }
1512}