tenant.openagents/omega

1use std::sync::Arc;
2
3use serde::{Deserialize, Serialize};
4
5use crate::role::Role;
6use crate::{
7    LanguageModelProviderId, LanguageModelToolUse, LanguageModelToolUseId,
8    LanguageModelToolUseInput, SharedString,
9};
10
11/// Dimensions of a `LanguageModelImage`
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub struct ImageSize {
14    pub width: i32,
15    pub height: i32,
16}
17
18#[derive(Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
19pub struct LanguageModelImage {
20    /// A base64-encoded PNG image.
21    pub source: SharedString,
22}
23
24impl LanguageModelImage {
25    pub fn len(&self) -> usize {
26        self.source.len()
27    }
28
29    pub fn is_empty(&self) -> bool {
30        self.source.is_empty()
31    }
32
33    pub fn empty() -> Self {
34        Self { source: "".into() }
35    }
36
37    /// Parse Self from a JSON object with case-insensitive field names
38    pub fn from_json(obj: &serde_json::Map<String, serde_json::Value>) -> Option<Self> {
39        let mut source = None;
40
41        for (k, v) in obj.iter() {
42            match k.to_lowercase().as_str() {
43                "source" => source = v.as_str(),
44                _ => {}
45            }
46        }
47
48        let source = source?;
49        Some(Self {
50            source: SharedString::from(source.to_string()),
51        })
52    }
53
54    pub fn to_base64_url(&self) -> String {
55        format!("data:image/png;base64,{}", self.source)
56    }
57}
58
59impl std::fmt::Debug for LanguageModelImage {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        f.debug_struct("LanguageModelImage")
62            .field("source", &format!("<{} bytes>", self.source.len()))
63            .finish()
64    }
65}
66
67#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
68pub struct LanguageModelToolResult {
69    pub tool_use_id: LanguageModelToolUseId,
70    pub tool_name: Arc<str>,
71    pub is_error: bool,
72    #[serde(with = "tool_result_content_vec")]
73    pub content: Vec<LanguageModelToolResultContent>,
74    /// The raw tool output, if available, often for debugging or extra state for replay
75    pub output: Option<serde_json::Value>,
76}
77
78impl LanguageModelToolResult {
79    /// Concatenates all `Text` parts of the content, ignoring non-text parts.
80    pub fn text_contents(&self) -> String {
81        let mut buffer = String::new();
82        for part in &self.content {
83            if let LanguageModelToolResultContent::Text(text) = part {
84                buffer.push_str(text);
85            }
86        }
87        buffer
88    }
89
90    /// Returns true when there are no content parts, or every part is empty.
91    pub fn is_content_empty(&self) -> bool {
92        self.content.iter().all(|part| part.is_empty())
93    }
94}
95
96/// Serde helper that accepts both the legacy single-value shape and the new
97/// array shape for `LanguageModelToolResult::content`, and normalizes both to
98/// `Vec<LanguageModelToolResultContent>`.
99mod tool_result_content_vec {
100    use super::LanguageModelToolResultContent;
101    use serde::{Deserialize, Deserializer, Serialize, Serializer};
102
103    pub fn serialize<S>(
104        value: &Vec<LanguageModelToolResultContent>,
105        serializer: S,
106    ) -> Result<S::Ok, S::Error>
107    where
108        S: Serializer,
109    {
110        value.serialize(serializer)
111    }
112
113    pub fn deserialize<'de, D>(
114        deserializer: D,
115    ) -> Result<Vec<LanguageModelToolResultContent>, D::Error>
116    where
117        D: Deserializer<'de>,
118    {
119        let value = serde_json::Value::deserialize(deserializer)?;
120        match value {
121            serde_json::Value::Array(items) => {
122                let mut out = Vec::with_capacity(items.len());
123                for item in items {
124                    out.push(
125                        serde_json::from_value::<LanguageModelToolResultContent>(item)
126                            .map_err(serde::de::Error::custom)?,
127                    );
128                }
129                Ok(out)
130            }
131            other => {
132                let single = serde_json::from_value::<LanguageModelToolResultContent>(other)
133                    .map_err(serde::de::Error::custom)?;
134                Ok(vec![single])
135            }
136        }
137    }
138}
139
140#[derive(Debug, Clone, Serialize, Eq, PartialEq, Hash)]
141pub enum LanguageModelToolResultContent {
142    Text(Arc<str>),
143    Image(LanguageModelImage),
144}
145
146impl<'de> Deserialize<'de> for LanguageModelToolResultContent {
147    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
148    where
149        D: serde::Deserializer<'de>,
150    {
151        use serde::de::Error;
152
153        let value = serde_json::Value::deserialize(deserializer)?;
154
155        // 1. Try as plain string
156        if let Ok(text) = serde_json::from_value::<String>(value.clone()) {
157            return Ok(Self::Text(Arc::from(text)));
158        }
159
160        // 2. Try as object
161        if let Some(obj) = value.as_object() {
162            fn get_field<'a>(
163                obj: &'a serde_json::Map<String, serde_json::Value>,
164                field: &str,
165            ) -> Option<&'a serde_json::Value> {
166                obj.iter()
167                    .find(|(k, _)| k.to_lowercase() == field.to_lowercase())
168                    .map(|(_, v)| v)
169            }
170
171            // Accept wrapped text format: { "type": "text", "text": "..." }
172            if let (Some(type_value), Some(text_value)) =
173                (get_field(obj, "type"), get_field(obj, "text"))
174                && let Some(type_str) = type_value.as_str()
175                && type_str.to_lowercase() == "text"
176                && let Some(text) = text_value.as_str()
177            {
178                return Ok(Self::Text(Arc::from(text)));
179            }
180
181            // Check for wrapped Text variant: { "text": "..." }
182            if let Some((_key, value)) = obj.iter().find(|(k, _)| k.to_lowercase() == "text")
183                && obj.len() == 1
184            {
185                if let Some(text) = value.as_str() {
186                    return Ok(Self::Text(Arc::from(text)));
187                }
188            }
189
190            // Check for wrapped Image variant: { "image": { "source": "...", "size": ... } }
191            if let Some((_key, value)) = obj.iter().find(|(k, _)| k.to_lowercase() == "image")
192                && obj.len() == 1
193            {
194                if let Some(image_obj) = value.as_object()
195                    && let Some(image) = LanguageModelImage::from_json(image_obj)
196                {
197                    return Ok(Self::Image(image));
198                }
199            }
200
201            // Try as direct Image
202            if let Some(image) = LanguageModelImage::from_json(obj) {
203                return Ok(Self::Image(image));
204            }
205        }
206
207        Err(D::Error::custom(format!(
208            "data did not match any variant of LanguageModelToolResultContent. Expected either a string, \
209             an object with 'type': 'text', a wrapped variant like {{\"Text\": \"...\"}}, or an image object. Got: {}",
210            serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string())
211        )))
212    }
213}
214
215impl LanguageModelToolResultContent {
216    pub fn to_str(&self) -> Option<&str> {
217        match self {
218            Self::Text(text) => Some(text),
219            Self::Image(_) => None,
220        }
221    }
222
223    pub fn is_empty(&self) -> bool {
224        match self {
225            Self::Text(text) => text.chars().all(|c| c.is_whitespace()),
226            Self::Image(_) => false,
227        }
228    }
229}
230
231impl From<&str> for LanguageModelToolResultContent {
232    fn from(value: &str) -> Self {
233        Self::Text(Arc::from(value))
234    }
235}
236
237impl From<String> for LanguageModelToolResultContent {
238    fn from(value: String) -> Self {
239        Self::Text(Arc::from(value))
240    }
241}
242
243impl From<anyhow::Error> for LanguageModelToolResultContent {
244    fn from(error: anyhow::Error) -> Self {
245        Self::Text(Arc::from(error.to_string()))
246    }
247}
248
249impl From<LanguageModelImage> for LanguageModelToolResultContent {
250    fn from(image: LanguageModelImage) -> Self {
251        Self::Image(image)
252    }
253}
254
255#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
256pub enum MessageContent {
257    Text(String),
258    Thinking {
259        text: String,
260        signature: Option<String>,
261    },
262    RedactedThinking(String),
263    Image(LanguageModelImage),
264    ToolUse(LanguageModelToolUse),
265    ToolResult(LanguageModelToolResult),
266    Compaction(CompactedContext),
267}
268
269#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
270pub enum CompactedContext {
271    Summary {
272        content: Arc<str>,
273        /// Opaque state the producing backend needs round-tripped alongside
274        /// the summary (e.g. Anthropic's `encrypted_content`). `None` when the
275        /// summary stands alone.
276        #[serde(default)]
277        provider_state: Option<ProviderCompactionState>,
278    },
279    ProviderState(ProviderCompactionState),
280}
281
282/// Opaque context produced by a provider's native compaction mechanism.
283///
284/// Only the provider identified by `provider_id` may interpret `payload`.
285/// `format` lets that provider evolve its representation without exposing it
286/// through the shared language model API.
287#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
288pub struct ProviderCompactionState {
289    provider_id: LanguageModelProviderId,
290    format: SharedString,
291    payload: Arc<str>,
292}
293
294impl ProviderCompactionState {
295    pub fn new(
296        provider_id: LanguageModelProviderId,
297        format: impl Into<SharedString>,
298        payload: impl Into<Arc<str>>,
299    ) -> Self {
300        Self {
301            provider_id,
302            format: format.into(),
303            payload: payload.into(),
304        }
305    }
306
307    pub fn provider_id(&self) -> &LanguageModelProviderId {
308        &self.provider_id
309    }
310
311    pub fn format(&self) -> &str {
312        &self.format
313    }
314
315    pub fn payload(&self) -> &str {
316        &self.payload
317    }
318}
319
320impl MessageContent {
321    pub fn is_empty(&self) -> bool {
322        match self {
323            MessageContent::Text(text) => text.chars().all(|c| c.is_whitespace()),
324            MessageContent::Thinking { text, .. } => text.chars().all(|c| c.is_whitespace()),
325            MessageContent::ToolResult(tool_result) => tool_result.is_content_empty(),
326            MessageContent::RedactedThinking(_)
327            | MessageContent::ToolUse(_)
328            | MessageContent::Image(_)
329            | MessageContent::Compaction(_) => false,
330        }
331    }
332}
333
334impl From<String> for MessageContent {
335    fn from(value: String) -> Self {
336        MessageContent::Text(value)
337    }
338}
339
340impl From<&str> for MessageContent {
341    fn from(value: &str) -> Self {
342        MessageContent::Text(value.to_string())
343    }
344}
345
346#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Hash)]
347pub struct LanguageModelRequestMessage {
348    pub role: Role,
349    pub content: Vec<MessageContent>,
350    pub cache: bool,
351    #[serde(default, skip_serializing_if = "Option::is_none")]
352    pub reasoning_details: Option<Arc<serde_json::Value>>,
353}
354
355impl LanguageModelRequestMessage {
356    pub fn string_contents(&self) -> String {
357        let mut buffer = String::new();
358        for content in &self.content {
359            match content {
360                MessageContent::Text(text) => {
361                    buffer.push_str(text);
362                }
363                MessageContent::Thinking { text, .. } => {
364                    buffer.push_str(text);
365                }
366                MessageContent::ToolResult(tool_result) => {
367                    for part in &tool_result.content {
368                        if let LanguageModelToolResultContent::Text(text) = part {
369                            buffer.push_str(text);
370                        }
371                    }
372                }
373                MessageContent::RedactedThinking(_)
374                | MessageContent::ToolUse(_)
375                | MessageContent::Image(_)
376                | MessageContent::Compaction(_) => {}
377            }
378        }
379        buffer
380    }
381
382    pub fn contents_empty(&self) -> bool {
383        self.content.iter().all(|content| content.is_empty())
384    }
385}
386
387#[derive(Debug, PartialEq, Hash, Clone, Serialize, Deserialize)]
388pub struct LanguageModelRequestTool {
389    pub name: String,
390    pub description: String,
391    pub input: LanguageModelRequestToolInput,
392}
393
394impl LanguageModelRequestTool {
395    pub fn function(
396        name: String,
397        description: String,
398        input_schema: serde_json::Value,
399        use_input_streaming: bool,
400    ) -> Self {
401        Self {
402            name,
403            description,
404            input: LanguageModelRequestToolInput::Function {
405                input_schema,
406                use_input_streaming,
407            },
408        }
409    }
410}
411
412#[derive(Debug, PartialEq, Hash, Clone, Serialize, Deserialize)]
413pub enum LanguageModelRequestToolInput {
414    Function {
415        input_schema: serde_json::Value,
416        use_input_streaming: bool,
417    },
418    Custom {
419        format: Option<LanguageModelCustomToolFormat>,
420    },
421}
422
423#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
424pub enum LanguageModelCustomToolFormat {
425    Text,
426    Grammar {
427        syntax: LanguageModelCustomToolGrammarSyntax,
428        definition: String,
429    },
430}
431
432#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, Serialize, Deserialize)]
433#[serde(rename_all = "lowercase")]
434pub enum LanguageModelCustomToolGrammarSyntax {
435    Lark,
436    Regex,
437}
438
439#[derive(Debug, PartialEq, Hash, Clone, Serialize, Deserialize)]
440pub enum LanguageModelToolChoice {
441    Auto,
442    Any,
443    None,
444}
445
446#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
447#[serde(rename_all = "snake_case")]
448pub enum CompletionIntent {
449    UserPrompt,
450    Subagent,
451    ToolResults,
452    ThreadSummarization,
453    ThreadContextSummarization,
454    CreateFile,
455    EditFile,
456    InlineAssist,
457    TerminalInlineAssist,
458    GenerateGitCommitMessage,
459}
460
461#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
462pub struct LanguageModelRequest {
463    pub thread_id: Option<String>,
464    pub prompt_id: Option<String>,
465    pub intent: Option<CompletionIntent>,
466    pub messages: Vec<LanguageModelRequestMessage>,
467    pub tools: Vec<LanguageModelRequestTool>,
468    pub tool_choice: Option<LanguageModelToolChoice>,
469    pub stop: Vec<String>,
470    pub temperature: Option<f32>,
471    pub thinking_allowed: bool,
472    pub thinking_effort: Option<String>,
473    pub speed: Option<Speed>,
474    #[serde(default, skip_serializing_if = "Option::is_none")]
475    pub compact_at_tokens: Option<u64>,
476}
477
478impl LanguageModelRequest {
479    pub fn contains_custom_tool_input(&self) -> bool {
480        self.tools
481            .iter()
482            .any(|tool| matches!(tool.input, LanguageModelRequestToolInput::Custom { .. }))
483            || self.messages.iter().any(|message| {
484                message.content.iter().any(|content| {
485                    matches!(
486                        content,
487                        MessageContent::ToolUse(LanguageModelToolUse {
488                            input: LanguageModelToolUseInput::Text(_),
489                            ..
490                        })
491                    )
492                })
493            })
494    }
495}
496
497#[derive(
498    Clone, Copy, Default, Debug, Serialize, Deserialize, PartialEq, Eq, schemars::JsonSchema,
499)]
500#[serde(rename_all = "snake_case")]
501pub enum Speed {
502    #[default]
503    Standard,
504    Fast,
505}
506
507impl Speed {
508    pub fn toggle(self) -> Self {
509        match self {
510            Speed::Standard => Speed::Fast,
511            Speed::Fast => Speed::Standard,
512        }
513    }
514}
515
516#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
517pub struct LanguageModelResponseMessage {
518    pub role: Option<Role>,
519    pub content: Option<String>,
520}
521
522#[cfg(test)]
523mod tests {
524    use super::*;
525
526    #[test]
527    fn test_language_model_tool_result_content_deserialization() {
528        // Test plain string
529        let json = serde_json::json!("hello world");
530        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
531        assert_eq!(
532            content,
533            LanguageModelToolResultContent::Text(Arc::from("hello world"))
534        );
535
536        // Test wrapped text format: { "type": "text", "text": "..." }
537        let json = serde_json::json!({"type": "text", "text": "hello"});
538        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
539        assert_eq!(
540            content,
541            LanguageModelToolResultContent::Text(Arc::from("hello"))
542        );
543
544        // Test single-field text object: { "text": "..." }
545        let json = serde_json::json!({"text": "hello"});
546        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
547        assert_eq!(
548            content,
549            LanguageModelToolResultContent::Text(Arc::from("hello"))
550        );
551
552        // Test case-insensitive type field
553        let json = serde_json::json!({"Type": "Text", "Text": "hello"});
554        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
555        assert_eq!(
556            content,
557            LanguageModelToolResultContent::Text(Arc::from("hello"))
558        );
559
560        // Test image object
561        let json = serde_json::json!({
562            "source": "base64encodedimagedata",
563        });
564        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
565        match content {
566            LanguageModelToolResultContent::Image(image) => {
567                assert_eq!(image.source.as_ref(), "base64encodedimagedata");
568            }
569            _ => panic!("Expected Image variant"),
570        }
571
572        // Test wrapped image: { "image": { "source": "...", "size": ... } }
573        let json = serde_json::json!({
574            "image": {
575                "source": "wrappedimagedata",
576            }
577        });
578        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
579        match content {
580            LanguageModelToolResultContent::Image(image) => {
581                assert_eq!(image.source.as_ref(), "wrappedimagedata");
582            }
583            _ => panic!("Expected Image variant"),
584        }
585
586        // Test case insensitive
587        let json = serde_json::json!({
588            "Source": "caseinsensitive",
589        });
590        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
591        match content {
592            LanguageModelToolResultContent::Image(image) => {
593                assert_eq!(image.source.as_ref(), "caseinsensitive");
594            }
595            _ => panic!("Expected Image variant"),
596        }
597
598        // Test direct image object
599        let json = serde_json::json!({
600            "source": "directimage",
601        });
602        let content: LanguageModelToolResultContent = serde_json::from_value(json).unwrap();
603        match content {
604            LanguageModelToolResultContent::Image(image) => {
605                assert_eq!(image.source.as_ref(), "directimage");
606            }
607            _ => panic!("Expected Image variant"),
608        }
609    }
610
611    #[test]
612    fn test_language_model_tool_result_content_vec_deserialization() {
613        // Legacy single-value shape is normalized to a Vec.
614        let json = serde_json::json!({
615            "tool_use_id": "abc",
616            "tool_name": "echo",
617            "is_error": false,
618            "content": "hello",
619            "output": null,
620        });
621        let result: LanguageModelToolResult = serde_json::from_value(json).unwrap();
622        assert_eq!(
623            result.content,
624            vec![LanguageModelToolResultContent::Text(Arc::from("hello"))]
625        );
626
627        // Legacy wrapped single-value shape also works.
628        let json = serde_json::json!({
629            "tool_use_id": "abc",
630            "tool_name": "echo",
631            "is_error": false,
632            "content": {"type": "text", "text": "hello"},
633            "output": null,
634        });
635        let result: LanguageModelToolResult = serde_json::from_value(json).unwrap();
636        assert_eq!(
637            result.content,
638            vec![LanguageModelToolResultContent::Text(Arc::from("hello"))]
639        );
640
641        // New array shape with text + image deserializes into a Vec.
642        let json = serde_json::json!({
643            "tool_use_id": "abc",
644            "tool_name": "echo",
645            "is_error": false,
646            "content": [
647                {"type": "text", "text": "foo"},
648                {"source": "data", "size": {"width": 1, "height": 2}}
649            ],
650            "output": null,
651        });
652        let result: LanguageModelToolResult = serde_json::from_value(json).unwrap();
653        assert_eq!(result.content.len(), 2);
654        assert_eq!(
655            result.content[0],
656            LanguageModelToolResultContent::Text(Arc::from("foo"))
657        );
658        match &result.content[1] {
659            LanguageModelToolResultContent::Image(image) => {
660                assert_eq!(image.source.as_ref(), "data");
661            }
662            _ => panic!("Expected Image variant"),
663        }
664
665        // Round-tripping preserves multi-part content.
666        let roundtripped: LanguageModelToolResult =
667            serde_json::from_value(serde_json::to_value(&result).unwrap()).unwrap();
668        assert_eq!(roundtripped, result);
669    }
670
671    #[test]
672    fn test_string_contents_includes_all_tool_result_text_parts() {
673        let tool_result = LanguageModelToolResult {
674            tool_use_id: LanguageModelToolUseId::from("id".to_string()),
675            tool_name: Arc::from("tool"),
676            is_error: false,
677            content: vec![
678                LanguageModelToolResultContent::Text(Arc::from("first ")),
679                LanguageModelToolResultContent::Image(LanguageModelImage::empty()),
680                LanguageModelToolResultContent::Text(Arc::from("second")),
681            ],
682            output: None,
683        };
684        let message = LanguageModelRequestMessage {
685            role: Role::User,
686            content: vec![
687                MessageContent::Text("prefix ".to_string()),
688                MessageContent::ToolResult(tool_result),
689                MessageContent::Text(" suffix".to_string()),
690            ],
691            cache: false,
692            reasoning_details: None,
693        };
694        assert_eq!(message.string_contents(), "prefix first second suffix");
695    }
696}
697