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 ToolCalls(Vec<ToolCall>),
47 Done(Box<LlmCompletionMetadata>),
49 Error(String),
51}
52
53#[derive(Debug, Clone)]
64pub struct DiscoveredModel {
65 pub model_id: String,
67 pub display_name: Option<String>,
69 pub created_at: Option<DateTime<Utc>>,
71 pub owned_by: Option<String>,
73 pub discovered_profile: Option<crate::llm_models::LlmModelProfile>,
76}
77
78#[derive(Debug, Clone, Default)]
86pub struct LlmCompletionMetadata {
87 pub total_tokens: Option<u32>,
89 pub prompt_tokens: Option<u32>,
91 pub completion_tokens: Option<u32>,
93 pub cache_read_tokens: Option<u32>,
95 pub cache_creation_tokens: Option<u32>,
97 pub model: Option<String>,
99 pub finish_reason: Option<String>,
101 pub retry_metadata: Option<crate::llm_retry::RetryMetadata>,
103 pub response_id: Option<String>,
106 pub phase: Option<String>,
110}
111
112#[async_trait]
116pub trait LlmDriver: Send + Sync {
117 async fn chat_completion_stream(
119 &self,
120 messages: Vec<LlmMessage>,
121 config: &LlmCallConfig,
122 ) -> Result<LlmResponseStream>;
123
124 async fn chat_completion(
126 &self,
127 messages: Vec<LlmMessage>,
128 config: &LlmCallConfig,
129 ) -> Result<LlmResponse> {
130 use futures::StreamExt;
131
132 let mut stream = self.chat_completion_stream(messages, config).await?;
133 let mut text = String::new();
134 let mut thinking = String::new();
135 let mut thinking_signature: Option<String> = None;
136 let mut tool_calls = Vec::new();
137 let mut metadata = LlmCompletionMetadata::default();
138
139 while let Some(event) = stream.next().await {
140 match event? {
141 LlmStreamEvent::TextDelta(delta) => text.push_str(&delta),
142 LlmStreamEvent::ThinkingDelta(delta) => thinking.push_str(&delta),
143 LlmStreamEvent::ThinkingSignature(sig) => thinking_signature = Some(sig),
144 LlmStreamEvent::ToolCalls(calls) => tool_calls = calls,
145 LlmStreamEvent::Done(meta) => metadata = *meta,
146 LlmStreamEvent::Error(err) => return Err(crate::error::AgentLoopError::llm(err)),
147 }
148 }
149
150 Ok(LlmResponse {
151 text,
152 thinking: if thinking.is_empty() {
153 None
154 } else {
155 Some(thinking)
156 },
157 thinking_signature,
158 tool_calls: if tool_calls.is_empty() {
159 None
160 } else {
161 Some(tool_calls)
162 },
163 metadata,
164 })
165 }
166
167 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
175 Ok(None)
177 }
178
179 fn supports_compact(&self) -> bool {
188 false
190 }
191
192 async fn compact(&self, _request: CompactRequest) -> Result<Option<CompactResponse>> {
212 Ok(None)
214 }
215}
216
217#[async_trait]
219impl LlmDriver for Box<dyn LlmDriver> {
220 async fn chat_completion_stream(
221 &self,
222 messages: Vec<LlmMessage>,
223 config: &LlmCallConfig,
224 ) -> Result<LlmResponseStream> {
225 (**self).chat_completion_stream(messages, config).await
226 }
227
228 async fn chat_completion(
229 &self,
230 messages: Vec<LlmMessage>,
231 config: &LlmCallConfig,
232 ) -> Result<LlmResponse> {
233 (**self).chat_completion(messages, config).await
234 }
235
236 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
237 (**self).list_models().await
238 }
239
240 fn supports_compact(&self) -> bool {
241 (**self).supports_compact()
242 }
243
244 async fn compact(&self, request: CompactRequest) -> Result<Option<CompactResponse>> {
245 (**self).compact(request).await
246 }
247}
248
249#[derive(Debug, Clone)]
255pub struct LlmMessage {
256 pub role: LlmMessageRole,
257 pub content: LlmMessageContent,
258 pub tool_calls: Option<Vec<ToolCall>>,
259 pub tool_call_id: Option<String>,
260 pub phase: Option<crate::message::ExecutionPhase>,
265 pub thinking: Option<String>,
268 pub thinking_signature: Option<String>,
271}
272
273impl LlmMessage {
274 pub fn text(role: LlmMessageRole, content: impl Into<String>) -> Self {
276 Self {
277 role,
278 content: LlmMessageContent::Text(content.into()),
279 tool_calls: None,
280 tool_call_id: None,
281 phase: None,
282 thinking: None,
283 thinking_signature: None,
284 }
285 }
286
287 pub fn parts(role: LlmMessageRole, parts: Vec<LlmContentPart>) -> Self {
289 Self {
290 role,
291 content: LlmMessageContent::Parts(parts),
292 tool_calls: None,
293 tool_call_id: None,
294 phase: None,
295 thinking: None,
296 thinking_signature: None,
297 }
298 }
299
300 pub fn content_as_text(&self) -> String {
302 self.content.to_text()
303 }
304
305 pub fn prepend_text_prefix(&mut self, prefix: &str) {
310 match &mut self.content {
311 LlmMessageContent::Text(text) => {
312 *text = format!("{}{}", prefix, text);
313 }
314 LlmMessageContent::Parts(parts) => {
315 for part in parts.iter_mut() {
316 if let LlmContentPart::Text { text } = part {
317 *text = format!("{}{}", prefix, text);
318 return;
319 }
320 }
321 parts.insert(
323 0,
324 LlmContentPart::Text {
325 text: prefix.to_string(),
326 },
327 );
328 }
329 }
330 }
331}
332
333#[derive(Debug, Clone)]
335pub enum LlmMessageContent {
336 Text(String),
338 Parts(Vec<LlmContentPart>),
340}
341
342impl LlmMessageContent {
343 pub fn to_text(&self) -> String {
345 match self {
346 LlmMessageContent::Text(s) => s.clone(),
347 LlmMessageContent::Parts(parts) => parts
348 .iter()
349 .filter_map(|p| match p {
350 LlmContentPart::Text { text } => Some(text.clone()),
351 _ => None,
352 })
353 .collect::<Vec<_>>()
354 .join(""),
355 }
356 }
357
358 pub fn is_text(&self) -> bool {
360 matches!(self, LlmMessageContent::Text(_))
361 }
362
363 pub fn is_parts(&self) -> bool {
365 matches!(self, LlmMessageContent::Parts(_))
366 }
367}
368
369impl From<String> for LlmMessageContent {
370 fn from(s: String) -> Self {
371 LlmMessageContent::Text(s)
372 }
373}
374
375impl From<&str> for LlmMessageContent {
376 fn from(s: &str) -> Self {
377 LlmMessageContent::Text(s.to_string())
378 }
379}
380
381#[derive(Debug, Clone)]
383pub enum LlmContentPart {
384 Text { text: String },
386 Image { url: String },
388 Audio { url: String },
390}
391
392impl LlmContentPart {
393 pub fn text(text: impl Into<String>) -> Self {
395 LlmContentPart::Text { text: text.into() }
396 }
397
398 pub fn image(url: impl Into<String>) -> Self {
400 LlmContentPart::Image { url: url.into() }
401 }
402
403 pub fn audio(url: impl Into<String>) -> Self {
405 LlmContentPart::Audio { url: url.into() }
406 }
407}
408
409#[derive(Debug, Clone, PartialEq, Eq)]
411pub enum LlmMessageRole {
412 System,
413 User,
414 Assistant,
415 Tool,
416}
417
418#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
428pub struct ToolSearchConfig {
429 pub enabled: bool,
431 pub threshold: usize,
434}
435
436#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
438#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
439#[serde(rename_all = "snake_case")]
440pub enum PromptCacheStrategy {
441 #[default]
443 Auto,
444}
445
446#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
451#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
452pub struct PromptCacheConfig {
453 pub enabled: bool,
455 #[serde(default)]
457 pub strategy: PromptCacheStrategy,
458 #[serde(default, skip_serializing_if = "Option::is_none")]
465 pub gemini_cached_content: Option<String>,
466}
467
468#[derive(Debug, Clone)]
470pub struct LlmCallConfig {
471 pub model: String,
472 pub temperature: Option<f32>,
473 pub max_tokens: Option<u32>,
474 pub tools: Vec<ToolDefinition>,
475 pub reasoning_effort: Option<String>,
477 pub metadata: HashMap<String, String>,
481 pub previous_response_id: Option<String>,
484 pub tool_search: Option<ToolSearchConfig>,
486 pub prompt_cache: Option<PromptCacheConfig>,
488}
489
490impl From<&RuntimeAgent> for LlmCallConfig {
491 fn from(runtime_agent: &RuntimeAgent) -> Self {
492 Self {
493 model: runtime_agent.model.clone(),
494 temperature: runtime_agent.temperature,
495 max_tokens: runtime_agent.max_tokens,
496 tools: runtime_agent.tools.clone(),
497 reasoning_effort: None, metadata: HashMap::new(), previous_response_id: None,
500 tool_search: runtime_agent.tool_search.clone(),
501 prompt_cache: runtime_agent.prompt_cache.clone(),
502 }
503 }
504}
505
506#[derive(Debug, Clone)]
508pub struct LlmResponse {
509 pub text: String,
510 pub thinking: Option<String>,
512 pub thinking_signature: Option<String>,
514 pub tool_calls: Option<Vec<ToolCall>>,
515 pub metadata: LlmCompletionMetadata,
516}
517
518pub struct LlmCallConfigBuilder {
537 config: LlmCallConfig,
538}
539
540impl LlmCallConfigBuilder {
541 pub fn from(runtime_agent: &RuntimeAgent) -> Self {
543 Self {
544 config: LlmCallConfig::from(runtime_agent),
545 }
546 }
547
548 pub fn reasoning_effort(mut self, effort: impl Into<String>) -> Self {
550 self.config.reasoning_effort = Some(effort.into());
551 self
552 }
553
554 pub fn model(mut self, model: impl Into<String>) -> Self {
556 self.config.model = model.into();
557 self
558 }
559
560 pub fn temperature(mut self, temp: f32) -> Self {
562 self.config.temperature = Some(temp);
563 self
564 }
565
566 pub fn max_tokens(mut self, tokens: u32) -> Self {
568 self.config.max_tokens = Some(tokens);
569 self
570 }
571
572 pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
574 self.config.tools = tools;
575 self
576 }
577
578 pub fn metadata(mut self, metadata: HashMap<String, String>) -> Self {
583 self.config.metadata = metadata;
584 self
585 }
586
587 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
589 self.config.metadata.insert(key.into(), value.into());
590 self
591 }
592
593 pub fn previous_response_id(mut self, id: Option<String>) -> Self {
595 self.config.previous_response_id = id;
596 self
597 }
598
599 pub fn tool_search(mut self, config: ToolSearchConfig) -> Self {
601 self.config.tool_search = Some(config);
602 self
603 }
604
605 pub fn prompt_cache(mut self, config: PromptCacheConfig) -> Self {
607 self.config.prompt_cache = Some(config);
608 self
609 }
610
611 pub fn build(self) -> LlmCallConfig {
613 self.config
614 }
615}
616
617impl From<&crate::message::Message> for LlmMessage {
622 fn from(msg: &crate::message::Message) -> Self {
628 let role = match msg.role {
629 crate::message::MessageRole::System => LlmMessageRole::System,
630 crate::message::MessageRole::User => LlmMessageRole::User,
631 crate::message::MessageRole::Agent => LlmMessageRole::Assistant,
632 crate::message::MessageRole::ToolResult => LlmMessageRole::Tool,
633 };
634
635 let tool_calls: Vec<ToolCall> = msg
637 .tool_calls()
638 .into_iter()
639 .map(|tc| ToolCall {
640 id: tc.id.clone(),
641 name: tc.name.clone(),
642 arguments: tc.arguments.clone(),
643 })
644 .collect();
645
646 LlmMessage {
647 role,
648 content: LlmMessageContent::Text(msg.content_to_llm_string()),
649 tool_calls: if tool_calls.is_empty() {
650 None
651 } else {
652 Some(tool_calls)
653 },
654 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
655 phase: msg.phase,
656 thinking: msg.thinking.clone(),
657 thinking_signature: msg.thinking_signature.clone(),
658 }
659 }
660}
661
662use crate::traits::ResolvedImage;
667use uuid::Uuid;
668
669impl LlmMessage {
670 pub fn from_message_with_images(
690 msg: &crate::message::Message,
691 resolved_images: &HashMap<Uuid, ResolvedImage>,
692 ) -> Self {
693 use crate::message::{ContentPart, MessageRole};
694
695 let role = match msg.role {
696 MessageRole::System => LlmMessageRole::System,
697 MessageRole::User => LlmMessageRole::User,
698 MessageRole::Agent => LlmMessageRole::Assistant,
699 MessageRole::ToolResult => LlmMessageRole::Tool,
700 };
701
702 let mut parts: Vec<LlmContentPart> = Vec::new();
704 let mut tool_calls: Vec<ToolCall> = Vec::new();
705
706 for part in &msg.content {
707 match part {
708 ContentPart::Text(t) => {
709 parts.push(LlmContentPart::Text {
710 text: t.text.clone(),
711 });
712 }
713 ContentPart::Image(img) => {
714 if let Some(url) = &img.url {
716 parts.push(LlmContentPart::Image { url: url.clone() });
717 } else if let (Some(base64), Some(media_type)) = (&img.base64, &img.media_type)
718 {
719 let data_url = format!("data:{};base64,{}", media_type, base64);
720 parts.push(LlmContentPart::Image { url: data_url });
721 }
722 }
723 ContentPart::ImageFile(img_file) => {
724 if let Some(resolved) = resolved_images.get(&img_file.image_id.uuid()) {
726 parts.push(LlmContentPart::Image {
727 url: resolved.to_data_url(),
728 });
729 } else {
730 parts.push(LlmContentPart::Text {
732 text: format!("[Image not found: {}]", img_file.image_id),
733 });
734 }
735 }
736 ContentPart::ToolCall(tc) => {
737 tool_calls.push(ToolCall {
739 id: tc.id.clone(),
740 name: tc.name.clone(),
741 arguments: tc.arguments.clone(),
742 });
743 }
744 ContentPart::ToolResult(tr) => {
745 let text = if let Some(err) = &tr.error {
747 format!("Tool error: {}", err)
748 } else if let Some(res) = &tr.result {
749 serde_json::to_string(res).unwrap_or_else(|_| "{}".to_string())
750 } else {
751 "{}".to_string()
752 };
753 let text = truncate_tool_result(text);
757 parts.push(LlmContentPart::Text { text });
758 }
759 }
760 }
761
762 let content = if parts.len() == 1 && matches!(&parts[0], LlmContentPart::Text { .. }) {
764 if let LlmContentPart::Text { text } = &parts[0] {
766 LlmMessageContent::Text(text.clone())
767 } else {
768 LlmMessageContent::Parts(parts)
769 }
770 } else if parts.is_empty() {
771 LlmMessageContent::Text(String::new())
773 } else {
774 LlmMessageContent::Parts(parts)
776 };
777
778 LlmMessage {
779 role,
780 content,
781 tool_calls: if tool_calls.is_empty() {
782 None
783 } else {
784 Some(tool_calls)
785 },
786 tool_call_id: msg.tool_call_id().map(|s| s.to_string()),
787 phase: msg.phase,
788 thinking: msg.thinking.clone(),
789 thinking_signature: msg.thinking_signature.clone(),
790 }
791 }
792
793 pub fn message_has_image_files(msg: &crate::message::Message) -> bool {
795 msg.content.iter().any(|p| p.is_image_file())
796 }
797
798 pub fn extract_image_file_ids(msg: &crate::message::Message) -> Vec<Uuid> {
800 msg.content
801 .iter()
802 .filter_map(|p| match p {
803 crate::message::ContentPart::ImageFile(f) => Some(f.image_id.uuid()),
804 _ => None,
805 })
806 .collect()
807 }
808}
809
810#[derive(Debug, Clone, PartialEq, Eq, Hash)]
816pub enum ProviderType {
817 OpenAI,
820 AzureOpenAI,
822 OpenAICompletions,
825 Anthropic,
826 Gemini,
828 LlmSim,
830}
831
832impl std::str::FromStr for ProviderType {
833 type Err = String;
834
835 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
836 match s.to_lowercase().as_str() {
837 "openai" => Ok(ProviderType::OpenAI),
838 "azure_openai" => Ok(ProviderType::AzureOpenAI),
839 "openai_completions" => Ok(ProviderType::OpenAICompletions),
840 "anthropic" => Ok(ProviderType::Anthropic),
841 "gemini" => Ok(ProviderType::Gemini),
842 "llmsim" => Ok(ProviderType::LlmSim),
843 _ => Err(format!("Unknown provider type: {}", s)),
844 }
845 }
846}
847
848impl std::fmt::Display for ProviderType {
849 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
850 match self {
851 ProviderType::OpenAI => write!(f, "openai"),
852 ProviderType::AzureOpenAI => write!(f, "azure_openai"),
853 ProviderType::OpenAICompletions => write!(f, "openai_completions"),
854 ProviderType::Anthropic => write!(f, "anthropic"),
855 ProviderType::Gemini => write!(f, "gemini"),
856 ProviderType::LlmSim => write!(f, "llmsim"),
857 }
858 }
859}
860
861#[derive(Debug, Clone)]
863pub struct ProviderConfig {
864 pub provider_type: ProviderType,
866 pub api_key: Option<String>,
868 pub base_url: Option<String>,
870}
871
872impl ProviderConfig {
873 pub fn new(provider_type: ProviderType) -> Self {
875 Self {
876 provider_type,
877 api_key: None,
878 base_url: None,
879 }
880 }
881
882 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
884 self.api_key = Some(api_key.into());
885 self
886 }
887
888 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
890 self.base_url = Some(base_url.into());
891 self
892 }
893}
894
895pub type BoxedLlmDriver = Box<dyn LlmDriver>;
897
898pub type DriverFactory = Arc<dyn Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync>;
906
907#[derive(Clone, Default)]
927pub struct DriverRegistry {
928 factories: HashMap<ProviderType, DriverFactory>,
929}
930
931impl DriverRegistry {
932 pub fn new() -> Self {
934 Self {
935 factories: HashMap::new(),
936 }
937 }
938
939 pub fn register<F>(&mut self, provider_type: ProviderType, factory: F)
941 where
942 F: Fn(&str, Option<&str>) -> BoxedLlmDriver + Send + Sync + 'static,
943 {
944 self.factories.insert(provider_type, Arc::new(factory));
945 }
946
947 pub fn create_driver(&self, config: &ProviderConfig) -> Result<BoxedLlmDriver> {
955 let api_key = if config.provider_type == ProviderType::LlmSim {
957 config.api_key.as_deref().unwrap_or("")
959 } else {
960 config.api_key.as_ref().ok_or_else(|| {
961 AgentLoopError::llm(
962 "API key is required. Configure the API key in provider settings.",
963 )
964 })?
965 };
966
967 let factory = self.factories.get(&config.provider_type).ok_or_else(|| {
969 AgentLoopError::driver_not_registered(config.provider_type.to_string())
970 })?;
971
972 Ok(factory(api_key, config.base_url.as_deref()))
974 }
975
976 pub fn has_driver(&self, provider_type: &ProviderType) -> bool {
978 self.factories.contains_key(provider_type)
979 }
980
981 pub fn registered_providers(&self) -> Vec<ProviderType> {
983 self.factories.keys().cloned().collect()
984 }
985}
986
987const MAX_TOOL_RESULT_BYTES: usize = 64 * 1024;
992
993const TRUNCATION_SUFFIX: &str =
994 "\n\n[Output truncated — exceeded 64 KiB limit. Try quiet flags, pipes, or redirect to file.]";
995
996fn truncate_tool_result(text: String) -> String {
997 if text.len() <= MAX_TOOL_RESULT_BYTES {
998 return text;
999 }
1000 let content_budget = MAX_TOOL_RESULT_BYTES.saturating_sub(TRUNCATION_SUFFIX.len());
1001 let mut end = content_budget;
1002 while end > 0 && !text.is_char_boundary(end) {
1003 end -= 1;
1004 }
1005 let mut truncated = text[..end].to_string();
1006 truncated.push_str(TRUNCATION_SUFFIX);
1007 truncated
1008}
1009
1010#[cfg(test)]
1015mod tests {
1016 use super::*;
1017
1018 #[test]
1019 fn test_llm_call_config_builder_from_runtime_agent() {
1020 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1021 let llm_config = LlmCallConfigBuilder::from(&runtime_agent).build();
1022
1023 assert_eq!(llm_config.model, "gpt-4o");
1024 assert!(llm_config.reasoning_effort.is_none());
1025 assert!(llm_config.temperature.is_none());
1026 assert!(llm_config.max_tokens.is_none());
1027 assert!(llm_config.tools.is_empty());
1028 assert!(llm_config.metadata.is_empty());
1029 }
1030
1031 #[test]
1032 fn test_llm_call_config_builder_with_metadata() {
1033 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1034 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1035 .with_metadata("session_id", "session_abc123")
1036 .with_metadata("agent_id", "agent_xyz789")
1037 .build();
1038
1039 assert_eq!(
1040 llm_config.metadata.get("session_id"),
1041 Some(&"session_abc123".to_string())
1042 );
1043 assert_eq!(
1044 llm_config.metadata.get("agent_id"),
1045 Some(&"agent_xyz789".to_string())
1046 );
1047 }
1048
1049 #[test]
1050 fn test_llm_call_config_builder_with_metadata_hashmap() {
1051 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1052 let mut metadata = HashMap::new();
1053 metadata.insert("key1".to_string(), "value1".to_string());
1054 metadata.insert("key2".to_string(), "value2".to_string());
1055
1056 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1057 .metadata(metadata)
1058 .build();
1059
1060 assert_eq!(llm_config.metadata.get("key1"), Some(&"value1".to_string()));
1061 assert_eq!(llm_config.metadata.get("key2"), Some(&"value2".to_string()));
1062 }
1063
1064 #[test]
1065 fn test_llm_call_config_builder_with_reasoning_effort() {
1066 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1067 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1068 .reasoning_effort("high")
1069 .build();
1070
1071 assert_eq!(llm_config.reasoning_effort, Some("high".to_string()));
1072 }
1073
1074 #[test]
1075 fn test_llm_call_config_builder_with_all_options() {
1076 let runtime_agent = RuntimeAgent::new("You are helpful", "gpt-4o");
1077 let llm_config = LlmCallConfigBuilder::from(&runtime_agent)
1078 .model("claude-3-opus")
1079 .reasoning_effort("medium")
1080 .temperature(0.7)
1081 .max_tokens(1000)
1082 .build();
1083
1084 assert_eq!(llm_config.model, "claude-3-opus");
1085 assert_eq!(llm_config.reasoning_effort, Some("medium".to_string()));
1086 assert_eq!(llm_config.temperature, Some(0.7));
1087 assert_eq!(llm_config.max_tokens, Some(1000));
1088 }
1089
1090 #[test]
1091 fn test_provider_type_parsing() {
1092 assert_eq!(
1093 "openai".parse::<ProviderType>().unwrap(),
1094 ProviderType::OpenAI
1095 );
1096 assert_eq!(
1097 "openai_completions".parse::<ProviderType>().unwrap(),
1098 ProviderType::OpenAICompletions
1099 );
1100 assert_eq!(
1101 "azure_openai".parse::<ProviderType>().unwrap(),
1102 ProviderType::AzureOpenAI
1103 );
1104 assert_eq!(
1105 "anthropic".parse::<ProviderType>().unwrap(),
1106 ProviderType::Anthropic
1107 );
1108 assert_eq!(
1109 "gemini".parse::<ProviderType>().unwrap(),
1110 ProviderType::Gemini
1111 );
1112 assert!("ollama".parse::<ProviderType>().is_err());
1114 assert!("custom".parse::<ProviderType>().is_err());
1115 }
1116
1117 #[test]
1118 fn test_provider_type_display() {
1119 assert_eq!(ProviderType::OpenAI.to_string(), "openai");
1120 assert_eq!(ProviderType::AzureOpenAI.to_string(), "azure_openai");
1121 assert_eq!(
1122 ProviderType::OpenAICompletions.to_string(),
1123 "openai_completions"
1124 );
1125 assert_eq!(ProviderType::Anthropic.to_string(), "anthropic");
1126 assert_eq!(ProviderType::Gemini.to_string(), "gemini");
1127 }
1128
1129 #[test]
1130 fn test_provider_config_builder() {
1131 let config = ProviderConfig::new(ProviderType::Anthropic)
1132 .with_api_key("test-key")
1133 .with_base_url("https://custom.api.com");
1134
1135 assert_eq!(config.provider_type, ProviderType::Anthropic);
1136 assert_eq!(config.api_key, Some("test-key".to_string()));
1137 assert_eq!(config.base_url, Some("https://custom.api.com".to_string()));
1138 }
1139
1140 #[test]
1141 fn test_driver_registry_requires_api_key() {
1142 let mut registry = DriverRegistry::new();
1144 registry.register(ProviderType::OpenAI, |_api_key, _base_url| {
1145 struct MockDriver;
1147 #[async_trait]
1148 impl LlmDriver for MockDriver {
1149 async fn chat_completion_stream(
1150 &self,
1151 _messages: Vec<LlmMessage>,
1152 _config: &LlmCallConfig,
1153 ) -> Result<LlmResponseStream> {
1154 unimplemented!()
1155 }
1156 }
1157 Box::new(MockDriver)
1158 });
1159
1160 let config = ProviderConfig::new(ProviderType::OpenAI);
1162 let result = registry.create_driver(&config);
1163 assert!(result.is_err());
1164
1165 let config_with_key = ProviderConfig::new(ProviderType::OpenAI).with_api_key("test-key");
1167 let result = registry.create_driver(&config_with_key);
1168 assert!(result.is_ok());
1169 }
1170
1171 #[test]
1172 fn test_driver_registry_returns_error_for_unregistered_provider() {
1173 let registry = DriverRegistry::new();
1174 let config = ProviderConfig::new(ProviderType::Anthropic).with_api_key("test-key");
1175
1176 let result = registry.create_driver(&config);
1177
1178 if let Err(AgentLoopError::DriverNotRegistered(provider)) = result {
1180 assert_eq!(provider, "anthropic");
1181 } else {
1182 panic!("Expected DriverNotRegistered error");
1183 }
1184 }
1185
1186 #[test]
1187 fn test_driver_registry_registration() {
1188 let mut registry = DriverRegistry::new();
1189
1190 assert!(!registry.has_driver(&ProviderType::OpenAI));
1191 assert!(!registry.has_driver(&ProviderType::Anthropic));
1192
1193 registry.register(ProviderType::OpenAI, |_, _| {
1194 struct MockDriver;
1195 #[async_trait]
1196 impl LlmDriver for MockDriver {
1197 async fn chat_completion_stream(
1198 &self,
1199 _messages: Vec<LlmMessage>,
1200 _config: &LlmCallConfig,
1201 ) -> Result<LlmResponseStream> {
1202 unimplemented!()
1203 }
1204 }
1205 Box::new(MockDriver)
1206 });
1207
1208 assert!(registry.has_driver(&ProviderType::OpenAI));
1209 assert!(!registry.has_driver(&ProviderType::Anthropic));
1210 }
1211
1212 use crate::{ContentPart, ImageFileContentPart, Message, MessageRole, TextContentPart};
1217
1218 #[test]
1219 fn test_message_has_image_files_with_image_file() {
1220 let message = Message {
1221 id: uuid::Uuid::new_v4().into(),
1222 role: MessageRole::User,
1223 content: vec![
1224 ContentPart::Text(TextContentPart {
1225 text: "Look at this image".to_string(),
1226 }),
1227 ContentPart::ImageFile(ImageFileContentPart {
1228 image_id: uuid::Uuid::new_v4().into(),
1229 filename: Some("test.png".to_string()),
1230 }),
1231 ],
1232 phase: None,
1233 thinking: None,
1234 thinking_signature: None,
1235 controls: None,
1236 metadata: None,
1237 external_actor: None,
1238 created_at: chrono::Utc::now(),
1239 };
1240
1241 assert!(LlmMessage::message_has_image_files(&message));
1242 }
1243
1244 #[test]
1245 fn test_message_has_image_files_without_image_file() {
1246 let message = Message {
1247 id: uuid::Uuid::new_v4().into(),
1248 role: MessageRole::User,
1249 content: vec![ContentPart::Text(TextContentPart {
1250 text: "Just text".to_string(),
1251 })],
1252 phase: None,
1253 thinking: None,
1254 thinking_signature: None,
1255 controls: None,
1256 metadata: None,
1257 external_actor: None,
1258 created_at: chrono::Utc::now(),
1259 };
1260
1261 assert!(!LlmMessage::message_has_image_files(&message));
1262 }
1263
1264 #[test]
1265 fn test_extract_image_file_ids() {
1266 let id1 = uuid::Uuid::new_v4();
1267 let id2 = uuid::Uuid::new_v4();
1268
1269 let message = Message {
1270 id: uuid::Uuid::new_v4().into(),
1271 role: MessageRole::User,
1272 content: vec![
1273 ContentPart::Text(TextContentPart {
1274 text: "Look at these images".to_string(),
1275 }),
1276 ContentPart::ImageFile(ImageFileContentPart {
1277 image_id: id1.into(),
1278 filename: Some("test1.png".to_string()),
1279 }),
1280 ContentPart::ImageFile(ImageFileContentPart {
1281 image_id: id2.into(),
1282 filename: Some("test2.png".to_string()),
1283 }),
1284 ],
1285 phase: None,
1286 thinking: None,
1287 thinking_signature: None,
1288 controls: None,
1289 metadata: None,
1290 external_actor: None,
1291 created_at: chrono::Utc::now(),
1292 };
1293
1294 let ids = LlmMessage::extract_image_file_ids(&message);
1295 assert_eq!(ids.len(), 2);
1296 assert!(ids.contains(&id1));
1297 assert!(ids.contains(&id2));
1298 }
1299
1300 #[test]
1301 fn test_from_message_with_images_text_only() {
1302 let message = Message {
1303 id: uuid::Uuid::new_v4().into(),
1304 role: MessageRole::User,
1305 content: vec![ContentPart::Text(TextContentPart {
1306 text: "Hello".to_string(),
1307 })],
1308 phase: None,
1309 thinking: None,
1310 thinking_signature: None,
1311 controls: None,
1312 metadata: None,
1313 external_actor: None,
1314 created_at: chrono::Utc::now(),
1315 };
1316
1317 let resolved = std::collections::HashMap::new();
1318 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1319
1320 assert_eq!(llm_message.role, LlmMessageRole::User);
1321 match llm_message.content {
1322 LlmMessageContent::Text(text) => assert_eq!(text, "Hello"),
1323 _ => panic!("Expected text content"),
1324 }
1325 }
1326
1327 #[test]
1328 fn test_from_message_with_images_resolved_image() {
1329 let image_id = uuid::Uuid::new_v4();
1330 let message = Message {
1331 id: uuid::Uuid::new_v4().into(),
1332 role: MessageRole::User,
1333 content: vec![
1334 ContentPart::Text(TextContentPart {
1335 text: "Look at this".to_string(),
1336 }),
1337 ContentPart::ImageFile(ImageFileContentPart {
1338 image_id: image_id.into(),
1339 filename: Some("test.png".to_string()),
1340 }),
1341 ],
1342 phase: None,
1343 thinking: None,
1344 thinking_signature: None,
1345 controls: None,
1346 metadata: None,
1347 external_actor: None,
1348 created_at: chrono::Utc::now(),
1349 };
1350
1351 let mut resolved = std::collections::HashMap::new();
1352 resolved.insert(
1353 image_id,
1354 crate::ResolvedImage::new("base64data", "image/png"),
1355 );
1356
1357 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1358
1359 match &llm_message.content {
1360 LlmMessageContent::Parts(parts) => {
1361 assert_eq!(parts.len(), 2);
1362 assert!(matches!(&parts[0], LlmContentPart::Text { .. }));
1364 if let LlmContentPart::Image { url } = &parts[1] {
1366 assert!(url.starts_with("data:image/png;base64,"));
1367 } else {
1368 panic!("Expected image content part");
1369 }
1370 }
1371 _ => panic!("Expected parts content"),
1372 }
1373 }
1374
1375 #[test]
1376 fn test_from_message_with_images_unresolved_image() {
1377 let image_id = uuid::Uuid::new_v4();
1378 let message = Message {
1379 id: uuid::Uuid::new_v4().into(),
1380 role: MessageRole::User,
1381 content: vec![ContentPart::ImageFile(ImageFileContentPart {
1382 image_id: image_id.into(),
1383 filename: Some("missing.png".to_string()),
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 resolved = std::collections::HashMap::new();
1396 let llm_message = LlmMessage::from_message_with_images(&message, &resolved);
1397
1398 match &llm_message.content {
1401 LlmMessageContent::Text(text) => {
1402 assert!(text.contains("Image not found"));
1403 }
1404 LlmMessageContent::Parts(parts) => {
1405 assert_eq!(parts.len(), 1);
1406 if let LlmContentPart::Text { text } = &parts[0] {
1407 assert!(text.contains("Image not found"));
1408 } else {
1409 panic!("Expected text placeholder for missing image");
1410 }
1411 }
1412 }
1413 }
1414
1415 #[test]
1416 fn test_prepend_text_prefix_simple_text() {
1417 let mut msg = LlmMessage::text(LlmMessageRole::User, "Hello bot");
1418 msg.prepend_text_prefix("[Alice] ");
1419 assert_eq!(msg.content_as_text(), "[Alice] Hello bot");
1420 }
1421
1422 #[test]
1423 fn test_prepend_text_prefix_parts() {
1424 let mut msg = LlmMessage::parts(
1425 LlmMessageRole::User,
1426 vec![
1427 LlmContentPart::Text {
1428 text: "Hello".to_string(),
1429 },
1430 LlmContentPart::Image {
1431 url: "data:image/png;base64,abc".to_string(),
1432 },
1433 ],
1434 );
1435 msg.prepend_text_prefix("[Bob] ");
1436 match &msg.content {
1437 LlmMessageContent::Parts(parts) => {
1438 if let LlmContentPart::Text { text } = &parts[0] {
1439 assert_eq!(text, "[Bob] Hello");
1440 } else {
1441 panic!("Expected text part");
1442 }
1443 }
1444 _ => panic!("Expected parts content"),
1445 }
1446 }
1447
1448 #[test]
1449 fn test_prepend_text_prefix_parts_no_text() {
1450 let mut msg = LlmMessage::parts(
1451 LlmMessageRole::User,
1452 vec![LlmContentPart::Image {
1453 url: "data:image/png;base64,abc".to_string(),
1454 }],
1455 );
1456 msg.prepend_text_prefix("[Eve] ");
1457 match &msg.content {
1458 LlmMessageContent::Parts(parts) => {
1459 assert_eq!(parts.len(), 2);
1460 if let LlmContentPart::Text { text } = &parts[0] {
1461 assert_eq!(text, "[Eve] ");
1462 } else {
1463 panic!("Expected prepended text part");
1464 }
1465 }
1466 _ => panic!("Expected parts content"),
1467 }
1468 }
1469}