1use serde::{Deserialize, Serialize};
2
3use super::GenerateModel;
4
5#[derive(Serialize, Default, Debug)]
6pub struct ChatRequest<'input> {
7 pub message: &'input str,
9 #[serde(skip_serializing_if = "Option::is_none")]
11 pub model: Option<GenerateModel>,
12 #[serde(skip_serializing_if = "Option::is_none")]
15 pub prompt_truncation: Option<PromptTruncation>,
16 #[serde(skip_serializing_if = "Option::is_none")]
18 pub temperature: Option<f64>,
19 #[serde(skip_serializing_if = "Option::is_none")]
22 pub conversation_id: Option<String>,
23 #[serde(skip_serializing_if = "Option::is_none")]
26 pub chat_history: Option<&'input Vec<ChatMessage>>,
27 #[serde(skip_serializing_if = "Option::is_none")]
29 pub preamble: Option<String>,
30 #[serde(skip_serializing_if = "Option::is_none")]
32 pub max_tokens: Option<u64>,
33 #[serde(skip_serializing_if = "Option::is_none")]
35 pub k: Option<u64>,
36}
37
38#[derive(Serialize, Deserialize, Debug, PartialEq)]
39#[serde(tag = "role")]
40pub enum ChatMessage {
41 #[serde(rename = "CHATBOT")]
42 Chatbot { message: String },
43 #[serde(rename = "USER")]
44 User { message: String },
45}
46
47#[derive(Serialize, Debug)]
48pub(crate) struct ChatStreamRequest<'input> {
49 #[serde(flatten)]
50 pub request: &'input ChatRequest<'input>,
51 pub stream: bool,
52}
53
54#[derive(strum_macros::Display, Serialize, Debug)]
55pub enum PromptTruncation {
56 #[strum(serialize = "AUTO")]
57 #[serde(rename = "AUTO")]
58 Auto,
59 #[strum(serialize = "OFF")]
60 #[serde(rename = "OFF")]
61 Off,
62}
63
64#[derive(strum_macros::Display, Serialize, Debug)]
65pub enum CitationQuality {
66 #[strum(serialize = "accurate")]
67 #[serde(rename = "accurate")]
68 Accurate,
69 #[strum(serialize = "fast")]
70 #[serde(rename = "fast")]
71 Fast,
72}
73
74#[derive(Deserialize, Debug, PartialEq)]
75pub struct ChatResponse {
76 pub generation_id: String,
77 pub response_id: String,
78 pub text: String,
79}
80
81#[derive(Deserialize, Debug, PartialEq)]
82#[serde(tag = "event_type")]
83pub enum ChatStreamResponse {
84 #[serde(rename = "stream-start")]
85 ChatStreamStart {
86 generation_id: String,
87 is_finished: bool,
88 },
89 #[serde(rename = "text-generation")]
90 ChatTextGeneration { is_finished: bool, text: String },
91 #[serde(rename = "stream-end")]
92 ChatStreamEnd {
93 finish_reason: String,
94 is_finished: bool,
95 response: ChatResponse,
96 },
97}