tiny_agent/transcript/
mod.rs1use crate::{error::StorageError, shared::Message};
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4
5#[derive(Deserialize, Serialize, Debug, Clone)]
6pub struct Transcript {
7 pub session_id: String,
8 pub messages: Vec<Message>,
9 pub parent: Option<String>,
11 pub subsessions: Vec<String>,
13}
14
15impl Transcript {
16 pub fn new(session_id: &str) -> Self {
17 Self {
18 session_id: String::from(session_id),
19 messages: Vec::new(),
20 parent: None,
21 subsessions: Vec::new(),
22 }
23 }
24
25 pub fn with_parent(session_id: &str, parent_session_id: &str) -> Self {
27 Self {
28 parent: Some(parent_session_id.to_string()),
29 ..Self::new(session_id)
30 }
31 }
32
33 pub fn append(&mut self, message: Message) {
34 self.messages.push(message);
35 }
36}
37
38#[async_trait]
39pub trait TranscriptStorage: Send + Sync {
40 async fn save_transcript(
41 &self,
42 session_id: &str,
43 transcript: &Transcript,
44 ) -> Result<(), StorageError>;
45 async fn get_transcript(&self, session_id: &str) -> Result<Option<Transcript>, StorageError>;
46 async fn append_message(&self, session_id: &str, message: Message) -> Result<(), StorageError> {
47 self.append_messages(session_id, vec![message]).await
48 }
49 async fn append_messages(
50 &self,
51 session_id: &str,
52 messages: Vec<Message>,
53 ) -> Result<(), StorageError> {
54 let mut transcript = self
55 .get_transcript(session_id)
56 .await?
57 .unwrap_or_else(|| Transcript::new(session_id));
58 for message in messages {
59 transcript.append(message);
60 }
61 self.save_transcript(session_id, &transcript).await
62 }
63
64 async fn add_subsession(
69 &self,
70 parent_session_id: &str,
71 child_session_id: &str,
72 ) -> Result<(), StorageError> {
73 let mut transcript = self
74 .get_transcript(parent_session_id)
75 .await?
76 .unwrap_or_else(|| Transcript::new(parent_session_id));
77 if !transcript
78 .subsessions
79 .iter()
80 .any(|s| s == child_session_id)
81 {
82 transcript.subsessions.push(child_session_id.to_string());
83 self.save_transcript(parent_session_id, &transcript).await?;
84 }
85 Ok(())
86 }
87
88 async fn get_subsessions(&self, session_id: &str) -> Result<Vec<String>, StorageError> {
90 Ok(self
91 .get_transcript(session_id)
92 .await?
93 .map(|t| t.subsessions)
94 .unwrap_or_default())
95 }
96
97 async fn remove_transcript(&self, session_id: &str) -> Result<(), StorageError>;
99
100 async fn delete_transcript(&self, session_id: &str) -> Result<(), StorageError> {
103 for child in self.get_subsessions(session_id).await? {
104 self.delete_transcript(&child).await?;
105 }
106 self.remove_transcript(session_id).await
107 }
108}