Skip to main content

aimux_core/
message.rs

1//! User-facing message types.
2//!
3//! `ModelMessage` is what the user passes to `generate_text` / `stream_text`.
4//! It is converted to `LanguageModelPrompt` (provider-facing) before calling
5//! `LanguageModel::do_generate` / `do_stream`.
6
7use serde::{Deserialize, Serialize};
8use ts_rs::TS;
9
10use crate::content::ContentPart;
11
12/// Who sent the message.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, TS)]
14#[serde(rename_all = "lowercase")]
15#[ts(export)]
16pub enum Role {
17    System,
18    /// The default role (a bare `Role::default()` yields `User`, matching the
19    /// common case for prompts built in tests).
20    #[default]
21    User,
22    Assistant,
23    Tool,
24}
25
26/// Message body: either a simple string or multi-part content.
27#[derive(Debug, Clone, Serialize, Deserialize, TS)]
28#[serde(untagged)]
29#[ts(export)]
30pub enum MessageContent {
31    Text(String),
32    Parts(Vec<ContentPart>),
33}
34
35/// A single user-facing chat message.
36#[derive(Debug, Clone, Serialize, Deserialize, TS)]
37#[ts(export)]
38pub struct ModelMessage {
39    pub role: Role,
40    pub content: MessageContent,
41}
42
43impl ModelMessage {
44    /// Convenience constructor for a text message.
45    pub fn text(role: Role, text: impl Into<String>) -> Self {
46        Self {
47            role,
48            content: MessageContent::Text(text.into()),
49        }
50    }
51
52    /// System message.
53    pub fn system(text: impl Into<String>) -> Self {
54        Self::text(Role::System, text)
55    }
56
57    /// User message.
58    pub fn user(text: impl Into<String>) -> Self {
59        Self::text(Role::User, text)
60    }
61
62    /// Assistant message.
63    pub fn assistant(text: impl Into<String>) -> Self {
64        Self::text(Role::Assistant, text)
65    }
66
67    /// User message with multi-part content.
68    #[must_use]
69    pub fn user_parts(parts: Vec<ContentPart>) -> Self {
70        Self {
71            role: Role::User,
72            content: MessageContent::Parts(parts),
73        }
74    }
75}
76
77/// What the user passes as `prompt` to `generate_text` / `stream_text`.
78///
79/// Can be a simple string (converted to a single user message) or a list of messages.
80#[derive(Debug, Clone, Serialize, Deserialize, TS)]
81#[serde(untagged)]
82#[ts(export)]
83pub enum ModelPrompt {
84    /// A plain string prompt — equivalent to a single user message.
85    Text(String),
86
87    /// A list of chat messages.
88    Messages(Vec<ModelMessage>),
89}
90
91impl From<String> for ModelPrompt {
92    fn from(s: String) -> Self {
93        ModelPrompt::Text(s)
94    }
95}
96
97impl From<&str> for ModelPrompt {
98    fn from(s: &str) -> Self {
99        ModelPrompt::Text(s.to_string())
100    }
101}
102
103impl From<Vec<ModelMessage>> for ModelPrompt {
104    fn from(msgs: Vec<ModelMessage>) -> Self {
105        ModelPrompt::Messages(msgs)
106    }
107}