Skip to main content

systemprompt_agent/services/a2a_server/processing/message/stream_processor/
mod.rs

1mod helpers;
2mod processing;
3
4use anyhow::{anyhow, Result};
5use base64::Engine;
6use std::sync::Arc;
7
8use crate::models::a2a::{FilePart, Message, Part};
9use crate::repository::execution::ExecutionStepRepository;
10use crate::services::{ContextService, SkillService};
11use systemprompt_models::{
12    is_supported_audio, is_supported_image, is_supported_text, is_supported_video, AiContentPart,
13    AiProvider,
14};
15
16#[allow(missing_debug_implementations)]
17pub struct StreamProcessor {
18    pub ai_service: Arc<dyn AiProvider>,
19    pub context_service: ContextService,
20    pub skill_service: Arc<SkillService>,
21    pub execution_step_repo: Arc<ExecutionStepRepository>,
22}
23
24impl StreamProcessor {
25    pub fn extract_message_text(message: &Message) -> Result<String> {
26        for part in &message.parts {
27            if let Part::Text(text_part) = part {
28                return Ok(text_part.text.clone());
29            }
30        }
31        Err(anyhow!("No text content found in message"))
32    }
33
34    pub fn extract_message_content(message: &Message) -> (String, Vec<AiContentPart>) {
35        let mut text_content = String::new();
36        let mut content_parts = Vec::new();
37
38        for part in &message.parts {
39            match part {
40                Part::Text(text_part) => {
41                    if text_content.is_empty() {
42                        text_content.clone_from(&text_part.text);
43                    }
44                    content_parts.push(AiContentPart::text(&text_part.text));
45                },
46                Part::File(file_part) => {
47                    if let Some(content_part) = Self::file_to_content_part(file_part) {
48                        content_parts.push(content_part);
49                    }
50                },
51                Part::Data(_) => {},
52            }
53        }
54
55        (text_content, content_parts)
56    }
57
58    fn file_to_content_part(file_part: &FilePart) -> Option<AiContentPart> {
59        let mime_type = file_part.file.mime_type.as_deref()?;
60        let file_name = file_part.file.name.as_deref().unwrap_or("unnamed");
61
62        if is_supported_image(mime_type) {
63            return Some(AiContentPart::image(mime_type, &file_part.file.bytes));
64        }
65
66        if is_supported_audio(mime_type) {
67            return Some(AiContentPart::audio(mime_type, &file_part.file.bytes));
68        }
69
70        if is_supported_video(mime_type) {
71            return Some(AiContentPart::video(mime_type, &file_part.file.bytes));
72        }
73
74        if is_supported_text(mime_type) {
75            return Self::decode_text_file(file_part, file_name, mime_type);
76        }
77
78        tracing::warn!(
79            file_name = %file_name,
80            mime_type = %mime_type,
81            "Unsupported file type - file will not be sent to AI"
82        );
83        None
84    }
85
86    fn decode_text_file(
87        file_part: &FilePart,
88        file_name: &str,
89        mime_type: &str,
90    ) -> Option<AiContentPart> {
91        let decoded = base64::engine::general_purpose::STANDARD
92            .decode(&file_part.file.bytes)
93            .map_err(|e| {
94                tracing::warn!(
95                    file_name = %file_name,
96                    mime_type = %mime_type,
97                    error = %e,
98                    "Failed to decode base64 text file"
99                );
100                e
101            })
102            .ok()?;
103
104        let text_content = String::from_utf8(decoded)
105            .map_err(|e| {
106                tracing::warn!(
107                    file_name = %file_name,
108                    mime_type = %mime_type,
109                    error = %e,
110                    "Failed to decode text file as UTF-8"
111                );
112                e
113            })
114            .ok()?;
115
116        let formatted = format!("[File: {file_name} ({mime_type})]\n{text_content}");
117        Some(AiContentPart::text(formatted))
118    }
119}