Skip to main content

systemprompt_runtime/trace/repository/
request.rs

1//! AI-request listing and detail queries.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use systemprompt_identifiers::{AiRequestId, TraceId, UserId};
7
8use super::{Result, TraceRepository};
9use crate::trace::models::{
10    AiRequestClientEvidence, AiRequestDetail, AiRequestFilter, AiRequestListItem,
11};
12
13struct DetailRow {
14    id: AiRequestId,
15    user_id: UserId,
16    actor_kind: String,
17    actor_id: String,
18    provider: Option<String>,
19    model: Option<String>,
20    input_tokens: Option<i32>,
21    output_tokens: Option<i32>,
22    cost_microdollars: i64,
23    latency_ms: Option<i32>,
24    status: String,
25    error_message: Option<String>,
26    client_kind: String,
27    client_attestation: String,
28    evidence_kind_source: Option<String>,
29    evidence_attested_host: Option<String>,
30    evidence_declared_client: Option<String>,
31    evidence_native_marker: Option<String>,
32    evidence_ua_product: Option<String>,
33    evidence_ua_version: Option<String>,
34    evidence_sdk_lang: Option<String>,
35    evidence_sdk_package_version: Option<String>,
36    evidence_sdk_runtime: Option<String>,
37    evidence_sdk_runtime_version: Option<String>,
38    evidence_sdk_os: Option<String>,
39    evidence_sdk_arch: Option<String>,
40}
41
42impl TraceRepository {
43    pub async fn list_ai_requests(
44        &self,
45        filter: &AiRequestFilter,
46    ) -> Result<Vec<AiRequestListItem>> {
47        let since = filter.since;
48        let until = filter.until;
49        let before_at = filter.before.as_ref().map(|c| c.created_at);
50        let before_id = filter.before.as_ref().map(|c| c.id.as_str());
51        let model = filter.model.as_deref();
52        let provider = filter.provider.as_deref();
53        let user = filter.user.as_deref();
54        let billing_source = filter.billing_source.as_deref();
55        let limit = filter.limit;
56        let rows = sqlx::query_as!(
57            AiRequestListItem,
58            r#"
59        SELECT
60            id as "id!: AiRequestId",
61            created_at as "created_at!",
62            trace_id as "trace_id: TraceId",
63            user_id as "user_id!: UserId",
64            actor_kind as "actor_kind!",
65            actor_id as "actor_id!",
66            client_kind as "client_kind!",
67            client_attestation as "client_attestation!",
68            provider, model,
69            input_tokens, output_tokens,
70            cache_read_tokens, cache_creation_tokens, reasoning_tokens,
71            cost_microdollars as "cost_microdollars!",
72            latency_ms,
73            status as "status!",
74            finish_reason,
75            billing_source as "billing_source!"
76        FROM ai_requests
77        WHERE ($1::timestamptz IS NULL OR created_at >= $1)
78          AND ($2::text IS NULL OR model ILIKE $2)
79          AND ($3::text IS NULL OR provider ILIKE $3)
80          AND ($4::text IS NULL OR user_id = $4)
81          AND ($6::timestamptz IS NULL OR created_at < $6)
82          AND ($7::timestamptz IS NULL OR (created_at, id) < ($7, $8))
83          AND ($9::text IS NULL OR billing_source = $9)
84        ORDER BY created_at DESC, id DESC
85        LIMIT $5
86        "#,
87            since,
88            model,
89            provider,
90            user,
91            limit,
92            until,
93            before_at,
94            before_id,
95            billing_source
96        )
97        .fetch_all(&*self.pool)
98        .await?;
99
100        Ok(rows)
101    }
102
103    pub async fn find_ai_request_detail(&self, id: &str) -> Result<Option<AiRequestDetail>> {
104        let partial = format!("{id}%");
105        let row = sqlx::query_as!(
106            DetailRow,
107            r#"
108        SELECT
109            r.id as "id!: AiRequestId",
110            r.user_id as "user_id!: UserId",
111            r.actor_kind as "actor_kind!",
112            r.actor_id as "actor_id!",
113            r.provider,
114            r.model,
115            r.input_tokens,
116            r.output_tokens,
117            r.cost_microdollars as "cost_microdollars!",
118            r.latency_ms,
119            r.status as "status!",
120            r.error_message,
121            r.client_kind as "client_kind!",
122            r.client_attestation as "client_attestation!",
123            e.kind_source as "evidence_kind_source?",
124            e.attested_host as "evidence_attested_host?",
125            e.declared_client as "evidence_declared_client?",
126            e.native_marker as "evidence_native_marker?",
127            e.ua_product as "evidence_ua_product?",
128            e.ua_version as "evidence_ua_version?",
129            e.sdk_lang as "evidence_sdk_lang?",
130            e.sdk_package_version as "evidence_sdk_package_version?",
131            e.sdk_runtime as "evidence_sdk_runtime?",
132            e.sdk_runtime_version as "evidence_sdk_runtime_version?",
133            e.sdk_os as "evidence_sdk_os?",
134            e.sdk_arch as "evidence_sdk_arch?"
135        FROM ai_requests r
136        LEFT JOIN ai_request_client_evidence e ON e.ai_request_id = r.id
137        WHERE r.id = $1 OR r.id LIKE $2
138        LIMIT 1
139        "#,
140            id,
141            partial
142        )
143        .fetch_optional(&*self.pool)
144        .await?;
145
146        Ok(row.map(|r| AiRequestDetail {
147            id: r.id,
148            user_id: r.user_id,
149            actor_kind: r.actor_kind,
150            actor_id: r.actor_id,
151            provider: r.provider,
152            model: r.model,
153            input_tokens: r.input_tokens,
154            output_tokens: r.output_tokens,
155            cost_microdollars: r.cost_microdollars,
156            latency_ms: r.latency_ms,
157            status: r.status,
158            error_message: r.error_message,
159            client_kind: r.client_kind,
160            client_attestation: r.client_attestation,
161            client_evidence: r
162                .evidence_kind_source
163                .map(|kind_source| AiRequestClientEvidence {
164                    kind_source,
165                    attested_host: r.evidence_attested_host,
166                    declared_client: r.evidence_declared_client,
167                    native_marker: r.evidence_native_marker,
168                    ua_product: r.evidence_ua_product,
169                    ua_version: r.evidence_ua_version,
170                    sdk_lang: r.evidence_sdk_lang,
171                    sdk_package_version: r.evidence_sdk_package_version,
172                    sdk_runtime: r.evidence_sdk_runtime,
173                    sdk_runtime_version: r.evidence_sdk_runtime_version,
174                    sdk_os: r.evidence_sdk_os,
175                    sdk_arch: r.evidence_sdk_arch,
176                }),
177        }))
178    }
179}