systemprompt_logging/trace/
ai_trace_service.rs

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

systemprompt_logging/trace/
ai_trace_service.rs
