anchor_chain/
message.rs

1//! Contains generic message types that can be used with multiple LLM models.
2//!
3//! These types are used to represent messages that can be sent to and received
4//! from LLM models. They are designed to be generic and reusable across
5//! different models and applications. This module also provides conversion
6//! traits to convert these generic message types into model-specific message
7//! types for use with specific LLM models.
8
9use async_openai::types::{
10    ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestMessage,
11    ChatCompletionRequestUserMessageArgs,
12};
13
14use crate::models::claude_3::{ClaudeMessage, ClaudeMessageContent};
15
16pub enum ChatMessage {
17    User(UserChatMessage),
18    Assistant(AssistantChatMessage),
19}
20
21#[allow(dead_code)]
22enum ChatMessageContent {
23    Text(String),
24    Image(String),
25}
26
27pub struct UserChatMessage {
28    content: ChatMessageContent,
29}
30
31pub struct AssistantChatMessage {
32    content: ChatMessageContent,
33}
34
35struct ChatMessageVec {
36    messages: Vec<ChatMessage>,
37}
38
39impl From<ChatMessage> for ChatCompletionRequestMessage {
40    fn from(message: ChatMessage) -> ChatCompletionRequestMessage {
41        match message {
42            ChatMessage::User(user_message) => user_message.into(),
43            ChatMessage::Assistant(assistant_message) => assistant_message.into(),
44        }
45    }
46}
47
48impl From<ChatMessage> for ClaudeMessage {
49    fn from(message: ChatMessage) -> ClaudeMessage {
50        match message {
51            ChatMessage::User(user_message) => user_message.into(),
52            ChatMessage::Assistant(assistant_message) => assistant_message.into(),
53        }
54    }
55}
56
57impl From<UserChatMessage> for ChatCompletionRequestMessage {
58    fn from(message: UserChatMessage) -> ChatCompletionRequestMessage {
59        match message.content {
60            ChatMessageContent::Text(text) => ChatCompletionRequestUserMessageArgs::default()
61                .content(text)
62                .build()
63                .unwrap()
64                .into(),
65            ChatMessageContent::Image(_image) => todo!("Image support"),
66        }
67    }
68}
69
70impl From<UserChatMessage> for ClaudeMessage {
71    fn from(message: UserChatMessage) -> ClaudeMessage {
72        match message.content {
73            ChatMessageContent::Text(text) => ClaudeMessage {
74                role: Some("user".to_string()),
75                content: vec![ClaudeMessageContent {
76                    content_type: "text".to_string(),
77                    text: Some(text),
78                    source: None,
79                }],
80            },
81            ChatMessageContent::Image(_image) => todo!("Image support"),
82        }
83    }
84}
85
86impl From<AssistantChatMessage> for ChatCompletionRequestMessage {
87    fn from(message: AssistantChatMessage) -> ChatCompletionRequestMessage {
88        match message.content {
89            ChatMessageContent::Text(text) => ChatCompletionRequestAssistantMessageArgs::default()
90                .content(text)
91                .build()
92                .unwrap()
93                .into(),
94            ChatMessageContent::Image(_image) => todo!("Image support"),
95        }
96    }
97}
98
99impl From<AssistantChatMessage> for ClaudeMessage {
100    fn from(message: AssistantChatMessage) -> ClaudeMessage {
101        match message.content {
102            ChatMessageContent::Text(text) => ClaudeMessage {
103                role: Some("assistant".to_string()),
104                content: vec![ClaudeMessageContent {
105                    content_type: "text".to_string(),
106                    text: Some(text),
107                    source: None,
108                }],
109            },
110            ChatMessageContent::Image(_image) => todo!("Image support"),
111        }
112    }
113}
114
115impl From<Vec<ChatMessage>> for ChatMessageVec {
116    fn from(messages: Vec<ChatMessage>) -> ChatMessageVec {
117        ChatMessageVec { messages }
118    }
119}
120
121impl From<ChatMessageVec> for Vec<ChatCompletionRequestMessage> {
122    fn from(messages: ChatMessageVec) -> Vec<ChatCompletionRequestMessage> {
123        messages.messages.into_iter().map(|m| m.into()).collect()
124    }
125}