diff options
| author | sigoden <sigoden@gmail.com> | 2023-11-02 10:45:11 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-11-02 10:45:11 +0800 |
| commit | 7c6841782d36faacc2aa3616dc3ed1b9403fa26e (patch) | |
| tree | 12d51728e056dbd4e57fc5977b7c51b2b60c06ae /src/client/message.rs | |
| parent | 444f4ebe9de8aa68afdf8d39b734a8543b9b5c4b (diff) | |
| download | aichat-7c6841782d36faacc2aa3616dc3ed1b9403fa26e.tar.gz | |
refactor: improve code quanity (#197)
- move model_info.rs/message.rs to clients/
- rename SharedConfig to GlobalConfig
Diffstat (limited to 'src/client/message.rs')
| -rw-r--r-- | src/client/message.rs | 52 |
1 files changed, 52 insertions, 0 deletions
diff --git a/src/client/message.rs b/src/client/message.rs new file mode 100644 index 0000000..55b2663 --- /dev/null +++ b/src/client/message.rs @@ -0,0 +1,52 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Message { + pub role: MessageRole, + pub content: String, +} + +impl Message { + pub fn new(content: &str) -> Self { + Self { + role: MessageRole::User, + content: content.to_string(), + } + } +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum MessageRole { + System, + Assistant, + User, +} + +#[allow(dead_code)] +impl MessageRole { + pub fn is_system(&self) -> bool { + matches!(self, MessageRole::System) + } + + pub fn is_user(&self) -> bool { + matches!(self, MessageRole::User) + } + + pub fn is_assistant(&self) -> bool { + matches!(self, MessageRole::Assistant) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_serde() { + assert_eq!( + serde_json::to_string(&Message::new("Hello World")).unwrap(), + "{\"role\":\"user\",\"content\":\"Hello World\"}" + ); + } +} |
