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

systemprompt_logging/trace/
ai_trace_service.rs
