systemprompt_runtime/trace/
ai_trace_service.rs

1//! Aggregated AI trace assembly across request, tool, and log rows.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use sqlx::PgPool;
7use std::sync::Arc;
8use systemprompt_identifiers::{AiRequestId, ContextId, McpExecutionId, TaskId};
9
10use super::models::{
11    AiRequestInfo, ConversationMessage, ExecutionStep, McpToolExecution, TaskArtifact, TaskInfo,
12    ToolLogEntry,
13};
14use super::{TraceError, ai_trace_queries};
15
16pub(super) type Result<T> = std::result::Result<T, TraceError>;
17
18#[derive(Debug, Clone)]
19pub struct AiTraceService {
20    pool: Arc<PgPool>,
21}
22
23impl AiTraceService {
24    pub const fn new(pool: Arc<PgPool>) -> Self {
25        Self { pool }
26    }
27
28    pub async fn resolve_task_id(&self, partial_id: &str) -> Result<TaskId> {
29        ai_trace_queries::resolve_task_id(&self.pool, partial_id)
30            .await?
31            .map(TaskId::new)
32            .ok_or_else(|| TraceError::TaskNotFound {
33                partial_id: partial_id.to_owned(),
34            })
35    }
36
37    pub async fn get_task_info(&self, task_id: &TaskId) -> Result<TaskInfo> {
38        ai_trace_queries::fetch_task_info(&self.pool, task_id).await
39    }
40
41    pub async fn get_user_input(&self, task_id: &TaskId) -> Result<Option<String>> {
42        ai_trace_queries::fetch_user_input(&self.pool, task_id).await
43    }
44
45    pub async fn get_agent_response(&self, task_id: &TaskId) -> Result<Option<String>> {
46        ai_trace_queries::fetch_agent_response(&self.pool, task_id).await
47    }
48
49    pub async fn get_execution_steps(&self, task_id: &TaskId) -> Result<Vec<ExecutionStep>> {
50        ai_trace_queries::fetch_execution_steps(&self.pool, task_id).await
51    }
52
53    pub async fn get_ai_requests(&self, task_id: &TaskId) -> Result<Vec<AiRequestInfo>> {
54        ai_trace_queries::fetch_ai_requests(&self.pool, task_id).await
55    }
56
57    pub async fn get_system_prompt(&self, request_id: &AiRequestId) -> Result<Option<String>> {
58        ai_trace_queries::fetch_system_prompt(&self.pool, request_id).await
59    }
60
61    pub async fn get_conversation_messages(
62        &self,
63        request_id: &AiRequestId,
64    ) -> Result<Vec<ConversationMessage>> {
65        ai_trace_queries::fetch_conversation_messages(&self.pool, request_id).await
66    }
67
68    pub async fn get_mcp_executions(
69        &self,
70        task_id: &TaskId,
71        context_id: &ContextId,
72    ) -> Result<Vec<McpToolExecution>> {
73        ai_trace_queries::fetch_mcp_executions(&self.pool, task_id, context_id).await
74    }
75
76    pub async fn get_mcp_linked_ai_requests(
77        &self,
78        mcp_execution_id: &McpExecutionId,
79    ) -> Result<Vec<AiRequestInfo>> {
80        ai_trace_queries::fetch_mcp_linked_ai_requests(&self.pool, mcp_execution_id).await
81    }
82
83    pub async fn get_ai_request_message_previews(
84        &self,
85        request_id: &AiRequestId,
86    ) -> Result<Vec<ConversationMessage>> {
87        ai_trace_queries::fetch_ai_request_message_previews(&self.pool, request_id).await
88    }
89
90    pub async fn get_tool_logs(
91        &self,
92        task_id: &TaskId,
93        context_id: &ContextId,
94    ) -> Result<Vec<ToolLogEntry>> {
95        ai_trace_queries::fetch_tool_logs(&self.pool, task_id, context_id).await
96    }
97
98    pub async fn get_task_artifacts(
99        &self,
100        task_id: &TaskId,
101        context_id: &ContextId,
102    ) -> Result<Vec<TaskArtifact>> {
103        ai_trace_queries::fetch_task_artifacts(&self.pool, task_id, context_id).await
104    }
105}
systemprompt_runtime/trace/ai_trace_service.rs

systemprompt_runtime/trace/
ai_trace_service.rs
