Skip to main content

baml_agent/session/
traits.rs

1//! Core traits and types for the session system.
2
3/// Role of a message in the agent conversation.
4pub trait MessageRole: Clone + PartialEq {
5    fn system() -> Self;
6    fn user() -> Self;
7    fn assistant() -> Self;
8    fn tool() -> Self;
9    fn as_str(&self) -> &str;
10    fn parse_role(s: &str) -> Option<Self>;
11    fn is_system(&self) -> bool {
12        self.as_str() == "system"
13    }
14}
15
16/// A message in the agent conversation.
17pub trait AgentMessage: Clone {
18    type Role: MessageRole;
19    fn new(role: Self::Role, content: String) -> Self;
20    fn role(&self) -> &Self::Role;
21    fn content(&self) -> &str;
22}
23
24/// Entry type discriminator — prevents invalid types at compile time.
25#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
26#[serde(rename_all = "lowercase")]
27pub enum EntryType {
28    User,
29    Assistant,
30    System,
31    Tool,
32}
33
34impl EntryType {
35    pub fn parse(s: &str) -> Option<Self> {
36        match s {
37            "user" => Some(Self::User),
38            "assistant" => Some(Self::Assistant),
39            "system" => Some(Self::System),
40            "tool" => Some(Self::Tool),
41            _ => None,
42        }
43    }
44
45    pub(crate) fn into_role<R: MessageRole>(self) -> R {
46        match self {
47            Self::User => R::user(),
48            Self::Assistant => R::assistant(),
49            Self::System => R::system(),
50            Self::Tool => R::tool(),
51        }
52    }
53}