systemprompt_logging/trace/
ai_trace_service.rs

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

systemprompt_logging/trace/
ai_trace_service.rs
