systemprompt_evaluation/repository/
sampling.rs

1//! Candidate selection over the `ai_requests` trace.
2//!
3//! Sampling excludes `actor_kind = 'job'` rows — judge and replay inference
4//! is attributed to a job actor, so without this exclusion each run would
5//! sample and grade the previous run's judge prompts.
6//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10use sqlx::PgPool;
11use std::sync::Arc;
12use systemprompt_database::DbPool;
13use systemprompt_identifiers::AiRequestId;
14
15use crate::error::Result;
16use crate::models::{CanonicalMessage, SampleFilter, SampledRequest};
17
18#[derive(Debug, Clone)]
19pub struct SamplingRepository {
20    pool: Arc<PgPool>,
21}
22
23impl SamplingRepository {
24    pub fn new(db: &DbPool) -> Result<Self> {
25        Ok(Self {
26            pool: db.pool_arc()?,
27        })
28    }
29
30    pub async fn sample(&self, filter: &SampleFilter) -> Result<Vec<SampledRequest>> {
31        let rows = sqlx::query!(
32            r#"
33            SELECT r.id, r.provider AS "provider!", r.model AS "model!",
34                   r.system_prompt_override, r.latency_ms, r.cost_microdollars,
35                   r.created_at, p.offered_tools, p.prepared_body_sha256
36            FROM ai_requests r
37            LEFT JOIN ai_request_payloads p ON p.ai_request_id = r.id
38            WHERE r.status = 'completed'
39              AND r.actor_kind <> 'job'
40              AND ($1::timestamptz IS NULL OR r.created_at >= $1)
41              AND ($2::timestamptz IS NULL OR r.created_at < $2)
42              AND ($3::text IS NULL OR r.provider = $3)
43              AND ($4::text IS NULL OR r.model = $4)
44              AND ($5::text[] IS NULL OR r.id = ANY($5))
45            ORDER BY r.created_at DESC
46            LIMIT $6
47            "#,
48            filter.since,
49            filter.until,
50            filter.provider.as_deref(),
51            filter.model.as_deref(),
52            filter.ids.as_deref(),
53            filter.limit
54        )
55        .fetch_all(self.pool.as_ref())
56        .await?;
57
58        let mut sampled = Vec::with_capacity(rows.len());
59        for row in rows {
60            let id = AiRequestId::new(row.id);
61            let (messages, response_text) = self.load_messages(&id).await?;
62            sampled.push(SampledRequest {
63                ai_request_id: id,
64                provider: row.provider,
65                model: row.model,
66                system_prompt_override: row.system_prompt_override,
67                messages,
68                response_text,
69                offered_tools: row.offered_tools,
70                prepared_body_sha256: row.prepared_body_sha256,
71                latency_ms: row.latency_ms,
72                cost_microdollars: row.cost_microdollars,
73                created_at: row.created_at,
74            });
75        }
76        Ok(sampled)
77    }
78
79    /// Cost as persisted by the audit path, looked up by the provider-facing
80    /// `request_id` (the UUID on `AiResponse`), not the row's primary key.
81    pub async fn request_cost(&self, request_id: &str) -> Result<i64> {
82        let cost = sqlx::query_scalar!(
83            "SELECT cost_microdollars FROM ai_requests WHERE request_id = $1",
84            request_id
85        )
86        .fetch_optional(self.pool.as_ref())
87        .await?;
88        Ok(cost.unwrap_or(0))
89    }
90
91    /// Splits the stored transcript into the prompt (everything up to the
92    /// last assistant message) and the response (that last assistant message).
93    async fn load_messages(
94        &self,
95        id: &AiRequestId,
96    ) -> Result<(Vec<CanonicalMessage>, Option<String>)> {
97        let rows = sqlx::query!(
98            r#"
99            SELECT role, content
100            FROM ai_request_messages
101            WHERE request_id = $1
102            ORDER BY sequence_number
103            "#,
104            id.as_str()
105        )
106        .fetch_all(self.pool.as_ref())
107        .await?;
108
109        let mut messages: Vec<CanonicalMessage> = rows
110            .into_iter()
111            .map(|row| CanonicalMessage {
112                role: row.role,
113                content: row.content,
114            })
115            .collect();
116
117        let response_text = match messages.last() {
118            Some(last) if last.role == "assistant" => messages.pop().map(|m| m.content),
119            _ => None,
120        };
121        Ok((messages, response_text))
122    }
123}
systemprompt_evaluation/repository/sampling.rs

systemprompt_evaluation/repository/
sampling.rs
