Skip to main content

systemprompt_agent/repository/context/message/
mod.rs

1mod parts;
2mod persistence;
3mod queries;
4
5use sqlx::PgPool;
6use std::sync::Arc;
7use systemprompt_database::DbPool;
8use systemprompt_identifiers::{ContextId, TaskId};
9use systemprompt_traits::RepositoryError;
10
11use crate::models::a2a::Message;
12
13pub use parts::{FileUploadContext, PersistPartSqlxParams, get_message_parts};
14pub use persistence::{
15    PersistMessageSqlxParams, PersistMessageWithTxParams, persist_message_sqlx,
16    persist_message_with_tx,
17};
18pub use queries::{
19    get_messages_by_context, get_messages_by_task, get_next_sequence_number,
20    get_next_sequence_number_in_tx, get_next_sequence_number_sqlx,
21};
22
23#[derive(Debug, Clone)]
24pub struct MessageRepository {
25    pool: Arc<PgPool>,
26}
27
28impl MessageRepository {
29    pub fn new(db: &DbPool) -> Result<Self, RepositoryError> {
30        let pool = db.pool_arc().map_err(|e| {
31            RepositoryError::InvalidData(format!("PostgreSQL pool not available: {e}"))
32        })?;
33        Ok(Self { pool })
34    }
35
36    pub async fn get_messages_by_task(
37        &self,
38        task_id: &TaskId,
39    ) -> Result<Vec<Message>, RepositoryError> {
40        get_messages_by_task(&self.pool, task_id).await
41    }
42
43    pub async fn get_messages_by_context(
44        &self,
45        context_id: &ContextId,
46    ) -> Result<Vec<Message>, RepositoryError> {
47        get_messages_by_context(&self.pool, context_id).await
48    }
49
50    pub async fn get_next_sequence_number(&self, task_id: &TaskId) -> Result<i32, RepositoryError> {
51        get_next_sequence_number(&self.pool, task_id).await
52    }
53
54    pub async fn persist_message_sqlx(
55        &self,
56        params: PersistMessageSqlxParams<'_>,
57    ) -> Result<(), RepositoryError> {
58        persist_message_sqlx(params).await
59    }
60
61    pub async fn persist_message_with_tx(
62        &self,
63        params: PersistMessageWithTxParams<'_>,
64    ) -> Result<(), RepositoryError> {
65        persist_message_with_tx(params).await
66    }
67}