systemprompt_runtime/trace/repository/
request.rs1use 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}