systemprompt_users/sessions/
ai_provider.rs1use async_trait::async_trait;
7use systemprompt_identifiers::SessionId;
8use systemprompt_traits::{
9 ActiveSession, AiProviderError, AiProviderResult, AiSessionProvider, CreateAiSessionParams,
10};
11
12use crate::repository::SessionRepository;
13use systemprompt_traits::session_store::{
14 CreateSessionParams, SessionAttribution, SessionIdentity, SessionNetwork,
15};
16
17#[derive(Debug)]
18pub struct UsersAiSessionProvider {
19 session_repo: SessionRepository,
20}
21
22impl UsersAiSessionProvider {
23 pub const fn from_repository(session_repo: SessionRepository) -> Self {
24 Self { session_repo }
25 }
26}
27
28#[async_trait]
29impl AiSessionProvider for UsersAiSessionProvider {
30 async fn create_session(&self, params: CreateAiSessionParams<'_>) -> AiProviderResult<()> {
31 let full_params = CreateSessionParams {
32 identity: SessionIdentity {
33 session_id: params.session_id,
34 user_id: params.user_id,
35 session_source: params.session_source,
36 fingerprint_hash: None,
37 },
38 network: SessionNetwork::default(),
39 attribution: SessionAttribution::default(),
40 is_bot: false,
41 is_ai_crawler: false,
42 expires_at: params.expires_at,
43 };
44
45 self.session_repo
46 .create_session(&full_params)
47 .await
48 .map_err(|e| AiProviderError::Internal(e.into()))
49 }
50
51 async fn increment_ai_usage(
52 &self,
53 session_id: &SessionId,
54 tokens: i32,
55 cost_microdollars: i64,
56 ) -> AiProviderResult<()> {
57 self.session_repo
58 .increment_ai_usage(session_id, tokens, cost_microdollars)
59 .await
60 .map_err(|e| AiProviderError::Internal(e.into()))
61 }
62
63 async fn find_live_session(
64 &self,
65 session_id: &SessionId,
66 ) -> AiProviderResult<Option<ActiveSession>> {
67 let session = self
68 .session_repo
69 .find_active_by_id(session_id)
70 .await
71 .map_err(|e| AiProviderError::Internal(e.into()))?;
72 Ok(session.map(|row| ActiveSession {
73 user_id: row.user_id,
74 }))
75 }
76}