Skip to main content

systemprompt_agent/services/
context.rs

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