systemprompt_runtime/trace/repository/
ai_trace.rs

1//! Queries assembling an AI task trace from the database.
2//!
3//! Fetches the components of a single agent task — task info, user input and
4//! agent response, execution steps, AI requests with their system prompt and
5//! conversation messages — and re-exports the MCP-execution query counterparts.
6//! `resolve_task_id` expands a partial task id to the most recent full match.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use systemprompt_identifiers::{AgentName, AiRequestId, ContextId, ExecutionStepId, TaskId};
12
13use super::{Result, TraceRepository};
14use crate::trace::models::{AiRequestInfo, ConversationMessage, ExecutionStep, TaskInfo};
15
16impl TraceRepository {
17    pub async fn resolve_task_id(&self, partial_id: &str) -> Result<Option<String>> {
18        let pattern = format!("{}%", partial_id);
19        let row = sqlx::query!(
20        "SELECT task_id FROM agent_tasks WHERE task_id LIKE $1 ORDER BY created_at DESC LIMIT 1",
21        pattern
22    )
23    .fetch_optional(&*self.pool)
24    .await?;
25
26        Ok(row.map(|r| r.task_id))
27    }
28
29    pub async fn fetch_task_info(&self, task_id: &TaskId) -> Result<TaskInfo> {
30        let row = sqlx::query!(
31            r#"SELECT task_id, context_id, agent_name, status, created_at, started_at, completed_at,
32                  execution_time_ms, error_message
33           FROM agent_tasks WHERE task_id = $1"#,
34            task_id.as_str()
35        )
36        .fetch_one(&*self.pool)
37        .await?;
38
39        Ok(TaskInfo {
40            task_id: TaskId::new(row.task_id),
41            context_id: ContextId::try_new(&row.context_id)?,
42            agent_name: row.agent_name.map(AgentName::new),
43            status: row.status,
44            created_at: row.created_at,
45            started_at: row.started_at,
46            completed_at: row.completed_at,
47            execution_time_ms: row.execution_time_ms,
48            error_message: row.error_message,
49        })
50    }
51
52    pub async fn fetch_user_input(&self, task_id: &TaskId) -> Result<Option<String>> {
53        let row = sqlx::query!(
54            r#"SELECT mp.text_content
55           FROM task_messages tm
56           JOIN message_parts mp ON mp.message_id = tm.message_id AND mp.task_id = tm.task_id
57           WHERE tm.task_id = $1 AND tm.role = 'user' AND mp.part_kind = 'text'
58           ORDER BY tm.sequence_number DESC LIMIT 1"#,
59            task_id.as_str()
60        )
61        .fetch_optional(&*self.pool)
62        .await?;
63
64        Ok(row.and_then(|r| r.text_content))
65    }
66
67    pub async fn fetch_agent_response(&self, task_id: &TaskId) -> Result<Option<String>> {
68        let row = sqlx::query!(
69            r#"SELECT mp.text_content
70           FROM task_messages tm
71           JOIN message_parts mp ON mp.message_id = tm.message_id AND mp.task_id = tm.task_id
72           WHERE tm.task_id = $1 AND tm.role = 'agent' AND mp.part_kind = 'text'
73           ORDER BY tm.sequence_number DESC LIMIT 1"#,
74            task_id.as_str()
75        )
76        .fetch_optional(&*self.pool)
77        .await?;
78
79        Ok(row.and_then(|r| r.text_content))
80    }
81
82    pub async fn fetch_execution_steps(&self, task_id: &TaskId) -> Result<Vec<ExecutionStep>> {
83        let rows = sqlx::query!(
84            r#"SELECT
85               step_id as id,
86               content->>'type' as step_type,
87               COALESCE(content->>'title', content->>'type') as title,
88               status,
89               duration_ms,
90               error_message
91           FROM task_execution_steps
92           WHERE task_id = $1
93           ORDER BY started_at"#,
94            task_id.as_str()
95        )
96        .fetch_all(&*self.pool)
97        .await?;
98
99        Ok(rows
100            .into_iter()
101            .map(|r| ExecutionStep {
102                step_id: ExecutionStepId::new(r.id),
103                step_type: r.step_type,
104                title: r.title,
105                status: r.status,
106                duration_ms: r.duration_ms,
107                error_message: r.error_message,
108            })
109            .collect())
110    }
111
112    pub async fn fetch_ai_requests(&self, task_id: &TaskId) -> Result<Vec<AiRequestInfo>> {
113        let rows = sqlx::query!(
114        r#"SELECT id, model, provider, max_tokens, input_tokens, output_tokens, cost_microdollars, latency_ms
115           FROM ai_requests
116           WHERE task_id = $1
117           ORDER BY created_at"#,
118        task_id.as_str()
119    )
120    .fetch_all(&*self.pool)
121    .await?;
122
123        Ok(rows
124            .into_iter()
125            .map(|r| AiRequestInfo {
126                id: AiRequestId::new(r.id),
127                provider: r.provider,
128                model: r.model,
129                max_tokens: r.max_tokens,
130                input_tokens: r.input_tokens,
131                output_tokens: r.output_tokens,
132                cost_microdollars: r.cost_microdollars,
133                latency_ms: r.latency_ms,
134            })
135            .collect())
136    }
137
138    pub async fn fetch_system_prompt(&self, request_id: &AiRequestId) -> Result<Option<String>> {
139        let row = sqlx::query!(
140            r#"SELECT content
141           FROM ai_request_messages
142           WHERE request_id = $1 AND role = 'system' AND sequence_number = 0
143           LIMIT 1"#,
144            request_id.as_str()
145        )
146        .fetch_optional(&*self.pool)
147        .await?;
148
149        Ok(row.map(|r| r.content))
150    }
151
152    pub async fn fetch_conversation_messages(
153        &self,
154        request_id: &AiRequestId,
155    ) -> Result<Vec<ConversationMessage>> {
156        let rows = sqlx::query!(
157            r#"SELECT role, content, sequence_number
158           FROM ai_request_messages
159           WHERE request_id = $1
160           ORDER BY sequence_number"#,
161            request_id.as_str()
162        )
163        .fetch_all(&*self.pool)
164        .await?;
165
166        Ok(rows
167            .into_iter()
168            .map(|r| ConversationMessage {
169                role: r.role,
170                content: r.content,
171                sequence_number: r.sequence_number,
172            })
173            .collect())
174    }
175
176    pub async fn fetch_ai_request_message_previews(
177        &self,
178        request_id: &AiRequestId,
179    ) -> Result<Vec<ConversationMessage>> {
180        let rows = sqlx::query!(
181            r#"SELECT role, LEFT(content, 500) as content_preview, sequence_number
182           FROM ai_request_messages
183           WHERE request_id = $1
184           ORDER BY sequence_number"#,
185            request_id.as_str()
186        )
187        .fetch_all(&*self.pool)
188        .await?;
189
190        Ok(rows
191            .into_iter()
192            .map(|r| ConversationMessage {
193                role: r.role,
194                content: r.content_preview.unwrap_or_else(String::new),
195                sequence_number: r.sequence_number,
196            })
197            .collect())
198    }
199}
systemprompt_runtime/trace/repository/ai_trace.rs

systemprompt_runtime/trace/repository/
ai_trace.rs
