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 limit = filter.limit;
55 let rows = sqlx::query!(
56 r#"
57 SELECT
58 id as "id!: AiRequestId",
59 created_at as "created_at!",
60 trace_id,
61 user_id as "user_id!: UserId",
62 actor_kind as "actor_kind!",
63 actor_id as "actor_id!",
64 client_kind as "client_kind!",
65 client_attestation as "client_attestation!",
66 provider, model,
67 input_tokens, output_tokens,
68 cache_read_tokens, cache_creation_tokens, reasoning_tokens,
69 cost_microdollars as "cost_microdollars!",
70 latency_ms,
71 status as "status!",
72 finish_reason
73 FROM ai_requests
74 WHERE ($1::timestamptz IS NULL OR created_at >= $1)
75 AND ($2::text IS NULL OR model ILIKE $2)
76 AND ($3::text IS NULL OR provider ILIKE $3)
77 AND ($4::text IS NULL OR user_id = $4)
78 AND ($6::timestamptz IS NULL OR created_at < $6)
79 AND ($7::timestamptz IS NULL OR (created_at, id) < ($7, $8))
80 ORDER BY created_at DESC, id DESC
81 LIMIT $5
82 "#,
83 since,
84 model,
85 provider,
86 user,
87 limit,
88 until,
89 before_at,
90 before_id
91 )
92 .fetch_all(&*self.pool)
93 .await?;
94
95 Ok(rows
96 .into_iter()
97 .map(|r| AiRequestListItem {
98 id: r.id,
99 created_at: r.created_at,
100 trace_id: r.trace_id.map(TraceId::new),
101 user_id: r.user_id,
102 actor_kind: r.actor_kind,
103 actor_id: r.actor_id,
104 client_kind: r.client_kind,
105 client_attestation: r.client_attestation,
106 provider: r.provider,
107 model: r.model,
108 input_tokens: r.input_tokens,
109 output_tokens: r.output_tokens,
110 cache_read_tokens: r.cache_read_tokens,
111 cache_creation_tokens: r.cache_creation_tokens,
112 reasoning_tokens: r.reasoning_tokens,
113 cost_microdollars: r.cost_microdollars,
114 latency_ms: r.latency_ms,
115 status: r.status,
116 finish_reason: r.finish_reason,
117 })
118 .collect())
119 }
120
121 pub async fn find_ai_request_detail(&self, id: &str) -> Result<Option<AiRequestDetail>> {
122 let partial = format!("{id}%");
123 let row = sqlx::query_as!(
124 DetailRow,
125 r#"
126 SELECT
127 r.id as "id!: AiRequestId",
128 r.user_id as "user_id!: UserId",
129 r.actor_kind as "actor_kind!",
130 r.actor_id as "actor_id!",
131 r.provider,
132 r.model,
133 r.input_tokens,
134 r.output_tokens,
135 r.cost_microdollars as "cost_microdollars!",
136 r.latency_ms,
137 r.status as "status!",
138 r.error_message,
139 r.client_kind as "client_kind!",
140 r.client_attestation as "client_attestation!",
141 e.kind_source as "evidence_kind_source?",
142 e.attested_host as "evidence_attested_host?",
143 e.declared_client as "evidence_declared_client?",
144 e.native_marker as "evidence_native_marker?",
145 e.ua_product as "evidence_ua_product?",
146 e.ua_version as "evidence_ua_version?",
147 e.sdk_lang as "evidence_sdk_lang?",
148 e.sdk_package_version as "evidence_sdk_package_version?",
149 e.sdk_runtime as "evidence_sdk_runtime?",
150 e.sdk_runtime_version as "evidence_sdk_runtime_version?",
151 e.sdk_os as "evidence_sdk_os?",
152 e.sdk_arch as "evidence_sdk_arch?"
153 FROM ai_requests r
154 LEFT JOIN ai_request_client_evidence e ON e.ai_request_id = r.id
155 WHERE r.id = $1 OR r.id LIKE $2
156 LIMIT 1
157 "#,
158 id,
159 partial
160 )
161 .fetch_optional(&*self.pool)
162 .await?;
163
164 Ok(row.map(|r| AiRequestDetail {
165 id: r.id,
166 user_id: r.user_id,
167 actor_kind: r.actor_kind,
168 actor_id: r.actor_id,
169 provider: r.provider,
170 model: r.model,
171 input_tokens: r.input_tokens,
172 output_tokens: r.output_tokens,
173 cost_microdollars: r.cost_microdollars,
174 latency_ms: r.latency_ms,
175 status: r.status,
176 error_message: r.error_message,
177 client_kind: r.client_kind,
178 client_attestation: r.client_attestation,
179 client_evidence: r
180 .evidence_kind_source
181 .map(|kind_source| AiRequestClientEvidence {
182 kind_source,
183 attested_host: r.evidence_attested_host,
184 declared_client: r.evidence_declared_client,
185 native_marker: r.evidence_native_marker,
186 ua_product: r.evidence_ua_product,
187 ua_version: r.evidence_ua_version,
188 sdk_lang: r.evidence_sdk_lang,
189 sdk_package_version: r.evidence_sdk_package_version,
190 sdk_runtime: r.evidence_sdk_runtime,
191 sdk_runtime_version: r.evidence_sdk_runtime_version,
192 sdk_os: r.evidence_sdk_os,
193 sdk_arch: r.evidence_sdk_arch,
194 }),
195 }))
196 }
197}