systemprompt_logging/trace/
ai_trace_service.rs

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

systemprompt_logging/trace/
ai_trace_service.rs
