Skip to main content

systemprompt_analytics/services/
ai_provider.rs

1use async_trait::async_trait;
2use systemprompt_database::DbPool;
3use systemprompt_identifiers::SessionId;
4use systemprompt_traits::{
5    AiProviderError, AiProviderResult, AiSessionProvider, CreateAiSessionParams,
6};
7
8use crate::repository::{CreateSessionParams, SessionRepository};
9
10#[derive(Debug)]
11pub struct AnalyticsAiSessionProvider {
12    session_repo: SessionRepository,
13}
14
15impl AnalyticsAiSessionProvider {
16    pub fn new(pool: &DbPool) -> anyhow::Result<Self> {
17        Ok(Self {
18            session_repo: SessionRepository::new(pool)?,
19        })
20    }
21
22    pub const fn from_repository(session_repo: SessionRepository) -> Self {
23        Self { session_repo }
24    }
25}
26
27#[async_trait]
28impl AiSessionProvider for AnalyticsAiSessionProvider {
29    async fn session_exists(&self, session_id: &SessionId) -> AiProviderResult<bool> {
30        self.session_repo
31            .exists(session_id)
32            .await
33            .map_err(|e| AiProviderError::Internal(e.to_string()))
34    }
35
36    async fn create_session(&self, params: CreateAiSessionParams<'_>) -> AiProviderResult<()> {
37        let full_params = CreateSessionParams {
38            session_id: params.session_id,
39            user_id: params.user_id,
40            session_source: params.session_source,
41            fingerprint_hash: None,
42            ip_address: None,
43            user_agent: None,
44            device_type: None,
45            browser: None,
46            os: None,
47            country: None,
48            region: None,
49            city: None,
50            preferred_locale: None,
51            referrer_source: None,
52            referrer_url: None,
53            landing_page: None,
54            entry_url: None,
55            utm_source: None,
56            utm_medium: None,
57            utm_campaign: None,
58            utm_content: None,
59            utm_term: None,
60            is_bot: false,
61            is_ai_crawler: false,
62            expires_at: params.expires_at,
63        };
64
65        self.session_repo
66            .create_session(&full_params)
67            .await
68            .map_err(|e| AiProviderError::Internal(e.to_string()))
69    }
70
71    async fn increment_ai_usage(
72        &self,
73        session_id: &SessionId,
74        tokens: i32,
75        cost_microdollars: i64,
76    ) -> AiProviderResult<()> {
77        self.session_repo
78            .increment_ai_usage(session_id, tokens, cost_microdollars)
79            .await
80            .map_err(|e| AiProviderError::Internal(e.to_string()))
81    }
82}