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

systemprompt_agent/services/
context.rs
