model_gateway_rs/model/
llm.rs

1use serde::{Deserialize, Serialize};
2
3use crate::model::role::Role;
4
5/// Single chat message.
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ChatMessage {
8    pub role: Role,
9    pub content: String,
10}
11
12impl ChatMessage {
13    pub fn user(content: &str) -> Self {
14        Self {
15            role: Role::User,
16            content: content.to_string(),
17        }
18    }
19    pub fn assistant(content: &str) -> Self {
20        Self {
21            role: Role::Assistant,
22            content: content.to_string(),
23        }
24    }
25    pub fn system(content: &str) -> Self {
26        Self {
27            role: Role::System,
28            content: content.to_string(),
29        }
30    }
31}
32
33pub struct ChatMessages(pub Vec<ChatMessage>);
34
35impl From<String> for ChatMessages {
36    fn from(content: String) -> Self {
37        ChatMessages(vec![ChatMessage::user(content.as_str())])
38    }
39}
40
41impl From<Vec<ChatMessage>> for ChatMessages {
42    fn from(messages: Vec<ChatMessage>) -> Self {
43        ChatMessages(messages)
44    }
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct LlmInput {
49    pub messages: Vec<ChatMessage>,
50    pub max_tokens: Option<u32>,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct LlmOutput {
55    pub message: Option<ChatMessage>,
56    pub usage: Option<u32>, // Placeholder for token usage
57}
58
59impl LlmOutput {
60    pub fn get_message(&self) -> Option<&ChatMessage> {
61        self.message.as_ref()
62    }
63
64    pub fn get_content(&self) -> &str {
65        match &self.message {
66            Some(msg) => &msg.content,
67            None => "",
68        }
69    }
70
71    pub fn get_usage(&self) -> Option<u32> {
72        self.usage
73    }
74}
model_gateway_rs/model/llm.rs

model_gateway_rs/model/
llm.rs
