systemprompt_agent/services/
context.rs

1use anyhow::{anyhow, Result};
2use base64::Engine;
3use systemprompt_database::DbPool;
4use systemprompt_models::{
5    is_supported_audio, is_supported_image, is_supported_text, is_supported_video, AiContentPart,
6    AiMessage, MessageRole,
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) -> Self {
19        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) = match Self::extract_message_content(&msg) {
38                        Ok((t, p)) if !t.is_empty() || !p.is_empty() => (t, p),
39                        Ok(_) => continue,
40                        Err(e) => {
41                            tracing::warn!(error = %e, "Failed to extract message content");
42                            continue;
43                        },
44                    };
45
46                    let role = match msg.role.as_str() {
47                        "user" => MessageRole::User,
48                        "agent" => MessageRole::Assistant,
49                        _ => continue,
50                    };
51
52                    history_messages.push(AiMessage {
53                        role,
54                        content: text,
55                        parts,
56                    });
57                }
58            }
59
60            if let Some(artifacts) = task.artifacts {
61                for artifact in artifacts {
62                    if let Ok(artifact_content) = Self::serialize_artifact_for_context(&artifact) {
63                        history_messages.push(AiMessage {
64                            role: MessageRole::Assistant,
65                            content: artifact_content,
66                            parts: Vec::new(),
67                        });
68                    }
69                }
70            }
71        }
72
73        Ok(history_messages)
74    }
75
76    fn extract_message_content(message: &Message) -> Result<(String, Vec<AiContentPart>)> {
77        let mut text_content = String::new();
78        let mut content_parts = Vec::new();
79
80        for part in &message.parts {
81            match part {
82                Part::Text(text_part) => {
83                    if text_content.is_empty() {
84                        text_content.clone_from(&text_part.text);
85                    }
86                    content_parts.push(AiContentPart::text(&text_part.text));
87                },
88                Part::File(file_part) => {
89                    if let Some(content_part) = Self::file_to_content_part(file_part) {
90                        content_parts.push(content_part);
91                    }
92                },
93                Part::Data(_) => {},
94            }
95        }
96
97        Ok((text_content, content_parts))
98    }
99
100    fn file_to_content_part(file_part: &FilePart) -> Option<AiContentPart> {
101        let mime_type = file_part.file.mime_type.as_deref()?;
102        let file_name = file_part.file.name.as_deref().unwrap_or("unnamed");
103
104        if is_supported_image(mime_type) {
105            return Some(AiContentPart::image(mime_type, &file_part.file.bytes));
106        }
107
108        if is_supported_audio(mime_type) {
109            return Some(AiContentPart::audio(mime_type, &file_part.file.bytes));
110        }
111
112        if is_supported_video(mime_type) {
113            return Some(AiContentPart::video(mime_type, &file_part.file.bytes));
114        }
115
116        if is_supported_text(mime_type) {
117            return Self::decode_text_file(file_part, file_name, mime_type);
118        }
119
120        tracing::warn!(
121            file_name = %file_name,
122            mime_type = %mime_type,
123            "Unsupported file type - file will not be sent to AI"
124        );
125        None
126    }
127
128    fn decode_text_file(
129        file_part: &FilePart,
130        file_name: &str,
131        mime_type: &str,
132    ) -> Option<AiContentPart> {
133        let decoded = base64::engine::general_purpose::STANDARD
134            .decode(&file_part.file.bytes)
135            .map_err(|e| {
136                tracing::warn!(
137                    file_name = %file_name,
138                    mime_type = %mime_type,
139                    error = %e,
140                    "Failed to decode base64 text file"
141                );
142                e
143            })
144            .ok()?;
145
146        let text_content = String::from_utf8(decoded)
147            .map_err(|e| {
148                tracing::warn!(
149                    file_name = %file_name,
150                    mime_type = %mime_type,
151                    error = %e,
152                    "Failed to decode text file as UTF-8"
153                );
154                e
155            })
156            .ok()?;
157
158        let formatted = format!("[File: {file_name} ({mime_type})]\n{text_content}");
159        Some(AiContentPart::text(formatted))
160    }
161
162    fn serialize_artifact_for_context(artifact: &Artifact) -> Result<String> {
163        let artifact_name = artifact
164            .name
165            .as_ref()
166            .map(String::as_str)
167            .unwrap_or("unnamed");
168
169        let mut content = format!(
170            "[Artifact: {} (type: {}, id: {})]",
171            artifact_name, artifact.metadata.artifact_type, artifact.id
172        );
173
174        if let Some(description) = &artifact.description {
175            if !description.is_empty() {
176                let truncated = if description.len() > 300 {
177                    format!("{}...", &description[..300])
178                } else {
179                    description.clone()
180                };
181                content.push_str(&format!("\n{truncated}"));
182            }
183        }
184
185        Ok(content)
186    }
187}
systemprompt_agent/services/context.rs

systemprompt_agent/services/
context.rs
