Skip to main content

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