model_gateway_rs/model/
llm.rs

1use serde::{Deserialize, Serialize};
2
3/// Role in chat messages.
4#[derive(Debug, Clone, Serialize, Deserialize)]
5#[serde(rename_all = "lowercase")]
6pub enum Role {
7    System,
8    User,
9    Assistant,
10}
11
12/// Single chat message.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct ChatMessage {
15    pub role: Role,
16    pub content: String,
17}
18
19impl ChatMessage {
20    pub fn user(content: &str) -> Self {
21        Self {
22            role: Role::User,
23            content: content.to_string(),
24        }
25    }
26    pub fn assistant(content: &str) -> Self {
27        Self {
28            role: Role::Assistant,
29            content: content.to_string(),
30        }
31    }
32    pub fn system(content: &str) -> Self {
33        Self {
34            role: Role::System,
35            content: content.to_string(),
36        }
37    }
38}
39
40pub struct ChatMessages(pub Vec<ChatMessage>);
41
42impl From<String> for ChatMessages {
43    fn from(content: String) -> Self {
44        ChatMessages(vec![ChatMessage::user(content.as_str())])
45    }
46}
47
48impl From<Vec<ChatMessage>> for ChatMessages {
49    fn from(messages: Vec<ChatMessage>) -> Self {
50        ChatMessages(messages)
51    }
52}
53
54/// Request body for chat completion.
55
56#[derive(Debug, Deserialize)]
57pub struct ChatChoice {
58    pub index: u32,
59    pub message: ChatMessage,
60    pub finish_reason: Option<String>,
61}
62
63#[derive(Debug, Deserialize)]
64pub struct ChatUsage {
65    pub prompt_tokens: u32,
66    pub completion_tokens: u32,
67    pub total_tokens: u32,
68}
69
70#[derive(Debug, Clone, Serialize)]
71pub struct ChatRequest {
72    pub model: String,
73    pub messages: Vec<ChatMessage>,
74    #[serde(skip_serializing_if = "Option::is_none")]
75    pub stream: Option<bool>,
76    #[serde(skip_serializing_if = "Option::is_none")]
77    pub temperature: Option<f32>,
78}
79
80#[derive(Debug, Deserialize)]
81pub struct ChatResponse {
82    pub id: String,
83    pub object: String,
84    pub created: u64,
85    pub model: String,
86    pub choices: Vec<ChatChoice>,
87    pub usage: Option<ChatUsage>,
88}
89
90impl ChatResponse {
91    /// Get the first choice's message content.
92    pub fn first_message(&self) -> Option<String> {
93        self.choices
94            .first()
95            .map(|choice| choice.message.content.clone())
96    }
97}
model_gateway_rs/model/llm.rs

model_gateway_rs/model/
llm.rs
