systemprompt_agent/services/
context.rs

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

systemprompt_agent/services/
context.rs
