systemprompt_agent/services/
context_provider.rs1use async_trait::async_trait;
6use systemprompt_database::DbPool;
7use systemprompt_identifiers::{ContextId, SessionId, UserId};
8use systemprompt_traits::{ContextProvider, ContextProviderError, ContextWithStats};
9
10use crate::error::AgentError;
11use crate::repository::ContextRepository;
12
13#[derive(Debug, Clone)]
14pub struct ContextProviderService {
15 repo: ContextRepository,
16}
17
18impl ContextProviderService {
19 pub fn new(db_pool: &DbPool) -> Result<Self, AgentError> {
20 Ok(Self {
21 repo: ContextRepository::new(db_pool)?,
22 })
23 }
24}
25
26#[async_trait]
27impl ContextProvider for ContextProviderService {
28 async fn list_contexts_with_stats(
29 &self,
30 user_id: &UserId,
31 ) -> Result<Vec<ContextWithStats>, ContextProviderError> {
32 let contexts = self
33 .repo
34 .list_contexts_with_stats(user_id)
35 .await
36 .map_err(|e| ContextProviderError::Database(e.to_string()))?;
37
38 Ok(contexts
39 .into_iter()
40 .map(|c| ContextWithStats {
41 context_id: c.context_id,
42 user_id: c.user_id,
43 name: c.name,
44 created_at: c.created_at,
45 updated_at: c.updated_at,
46 task_count: c.task_count,
47 message_count: c.message_count,
48 last_message_at: c.last_message_at,
49 })
50 .collect())
51 }
52
53 async fn get_context(
54 &self,
55 context_id: &ContextId,
56 user_id: &UserId,
57 ) -> Result<ContextWithStats, ContextProviderError> {
58 let context = self
59 .repo
60 .get_context(context_id, user_id)
61 .await
62 .map_err(|e| match e {
63 systemprompt_traits::RepositoryError::NotFound(msg) => {
64 ContextProviderError::NotFound(msg)
65 },
66 other => ContextProviderError::Database(other.to_string()),
67 })?;
68
69 let all_contexts = self
70 .repo
71 .list_contexts_with_stats(user_id)
72 .await
73 .map_err(|e| ContextProviderError::Database(e.to_string()))?;
74
75 let context_with_stats = all_contexts
76 .into_iter()
77 .find(|c| c.context_id == context.context_id)
78 .ok_or_else(|| {
79 ContextProviderError::NotFound(format!("Context {} not found", context_id))
80 })?;
81
82 Ok(ContextWithStats {
83 context_id: context_with_stats.context_id,
84 user_id: context_with_stats.user_id,
85 name: context_with_stats.name,
86 created_at: context_with_stats.created_at,
87 updated_at: context_with_stats.updated_at,
88 task_count: context_with_stats.task_count,
89 message_count: context_with_stats.message_count,
90 last_message_at: context_with_stats.last_message_at,
91 })
92 }
93
94 async fn create_context(
95 &self,
96 user_id: &UserId,
97 session_id: Option<&SessionId>,
98 name: &str,
99 ) -> Result<ContextId, ContextProviderError> {
100 self.repo
101 .create_context(user_id, session_id, name)
102 .await
103 .map_err(|e| ContextProviderError::Database(e.to_string()))
104 }
105
106 async fn update_context_name(
107 &self,
108 context_id: &ContextId,
109 user_id: &UserId,
110 name: &str,
111 ) -> Result<(), ContextProviderError> {
112 self.repo
113 .update_context_name(context_id, user_id, name)
114 .await
115 .map_err(|e| match e {
116 systemprompt_traits::RepositoryError::NotFound(msg) => {
117 ContextProviderError::NotFound(msg)
118 },
119 other => ContextProviderError::Database(other.to_string()),
120 })
121 }
122
123 async fn delete_context(
124 &self,
125 context_id: &ContextId,
126 user_id: &UserId,
127 ) -> Result<(), ContextProviderError> {
128 self.repo
129 .delete_context(context_id, user_id)
130 .await
131 .map_err(|e| match e {
132 systemprompt_traits::RepositoryError::NotFound(msg) => {
133 ContextProviderError::NotFound(msg)
134 },
135 other => ContextProviderError::Database(other.to_string()),
136 })
137 }
138}