Skip to main content

systemprompt_agent/services/a2a_server/processing/
conversation_service.rs

1use anyhow::{Result, anyhow};
2use base64::Engine;
3use systemprompt_database::DbPool;
4use systemprompt_identifiers::ContextId;
5use systemprompt_models::{
6    AiContentPart, AiMessage, MessageRole, is_supported_audio, is_supported_image,
7    is_supported_text, is_supported_video,
8};
9
10use crate::models::a2a::{FilePart, Part};
11use crate::models::{Artifact, Message};
12use crate::repository::task::TaskRepository;
13
14#[derive(Debug)]
15pub struct ConversationService {
16    db_pool: DbPool,
17}
18
19impl ConversationService {
20    pub const fn new(db_pool: DbPool) -> Self {
21        Self { db_pool }
22    }
23
24    pub async fn load_conversation_history(
25        &self,
26        context_id: &ContextId,
27    ) -> Result<Vec<AiMessage>> {
28        let task_repo = TaskRepository::new(&self.db_pool)?;
29        let tasks = task_repo
30            .list_tasks_by_context(context_id)
31            .await
32            .map_err(|e| anyhow!("Failed to load conversation history: {}", e))?;
33
34        let mut history_messages = Vec::new();
35
36        for task in tasks {
37            if let Some(task_history) = task.history {
38                for msg in task_history {
39                    let (text, parts) = Self::extract_message_content(&msg);
40                    if text.is_empty() && parts.is_empty() {
41                        continue;
42                    }
43
44                    let role = match msg.role.as_str() {
45                        "user" => MessageRole::User,
46                        "agent" => MessageRole::Assistant,
47                        _ => continue,
48                    };
49
50                    history_messages.push(AiMessage {
51                        role,
52                        content: text,
53                        parts,
54                    });
55                }
56            }
57
58            if let Some(artifacts) = task.artifacts {
59                for artifact in artifacts {
60                    let artifact_content = Self::serialize_artifact_for_context(&artifact);
61                    history_messages.push(AiMessage {
62                        role: MessageRole::Assistant,
63                        content: artifact_content,
64                        parts: Vec::new(),
65                    });
66                }
67            }
68        }
69
70        Ok(history_messages)
71    }
72
73    fn extract_message_content(message: &Message) -> (String, Vec<AiContentPart>) {
74        let mut text_content = String::new();
75        let mut content_parts = Vec::new();
76
77        for part in &message.parts {
78            match part {
79                Part::Text(text_part) => {
80                    if text_content.is_empty() {
81                        text_content.clone_from(&text_part.text);
82                    }
83                    content_parts.push(AiContentPart::text(&text_part.text));
84                },
85                Part::File(file_part) => {
86                    if let Some(content_part) = Self::file_to_content_part(file_part) {
87                        content_parts.push(content_part);
88                    }
89                },
90                Part::Data(_) => {},
91            }
92        }
93
94        (text_content, content_parts)
95    }
96
97    fn file_to_content_part(file_part: &FilePart) -> Option<AiContentPart> {
98        let mime_type = file_part.file.mime_type.as_deref()?;
99        let file_name = file_part.file.name.as_deref().unwrap_or("unnamed");
100
101        if is_supported_image(mime_type) {
102            return Some(AiContentPart::image(mime_type, &file_part.file.bytes));
103        }
104
105        if is_supported_audio(mime_type) {
106            return Some(AiContentPart::audio(mime_type, &file_part.file.bytes));
107        }
108
109        if is_supported_video(mime_type) {
110            return Some(AiContentPart::video(mime_type, &file_part.file.bytes));
111        }
112
113        if is_supported_text(mime_type) {
114            return Self::decode_text_file(file_part, file_name, mime_type);
115        }
116
117        tracing::warn!(
118            file_name = %file_name,
119            mime_type = %mime_type,
120            "Unsupported file type - file will not be sent to AI"
121        );
122        None
123    }
124
125    fn decode_text_file(
126        file_part: &FilePart,
127        file_name: &str,
128        mime_type: &str,
129    ) -> Option<AiContentPart> {
130        let decoded = base64::engine::general_purpose::STANDARD
131            .decode(&file_part.file.bytes)
132            .map_err(|e| {
133                tracing::warn!(
134                    file_name = %file_name,
135                    mime_type = %mime_type,
136                    error = %e,
137                    "Failed to decode base64 text file"
138                );
139                e
140            })
141            .ok()?;
142
143        let text_content = String::from_utf8(decoded)
144            .map_err(|e| {
145                tracing::warn!(
146                    file_name = %file_name,
147                    mime_type = %mime_type,
148                    error = %e,
149                    "Failed to decode text file as UTF-8"
150                );
151                e
152            })
153            .ok()?;
154
155        let formatted = format!("[File: {file_name} ({mime_type})]\n{text_content}");
156        Some(AiContentPart::text(formatted))
157    }
158
159    fn serialize_artifact_for_context(artifact: &Artifact) -> String {
160        let artifact_name = artifact
161            .name
162            .clone()
163            .unwrap_or_else(|| "unnamed".to_string());
164
165        let mut content = format!(
166            "[Artifact: {} (type: {})]\n",
167            artifact_name, artifact.metadata.artifact_type
168        );
169
170        for part in &artifact.parts {
171            match part {
172                Part::Text(text_part) => {
173                    content.push_str(&text_part.text);
174                    content.push('\n');
175                },
176                Part::Data(data_part) => {
177                    let json_str = serde_json::to_string_pretty(&data_part.data)
178                        .unwrap_or_else(|_| "{}".to_string());
179                    content.push_str(&json_str);
180                    content.push('\n');
181                },
182                Part::File(file_part) => {
183                    if let Some(name) = &file_part.file.name {
184                        content.push_str(&format!("[File: {}]\n", name));
185                    }
186                },
187            }
188        }
189
190        content
191    }
192}