1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum ChatRole {
18 System,
19 User,
20 Assistant,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct ChatMessage {
26 pub role: ChatRole,
27 pub content: String,
28}
29
30impl ChatMessage {
31 pub fn system(content: impl Into<String>) -> Self {
32 Self {
33 role: ChatRole::System,
34 content: content.into(),
35 }
36 }
37 pub fn user(content: impl Into<String>) -> Self {
38 Self {
39 role: ChatRole::User,
40 content: content.into(),
41 }
42 }
43 pub fn assistant(content: impl Into<String>) -> Self {
44 Self {
45 role: ChatRole::Assistant,
46 content: content.into(),
47 }
48 }
49}
50
51#[derive(Clone, PartialEq, Eq)]
60pub struct CredentialRef(String);
61
62impl CredentialRef {
63 pub fn new(key: impl Into<String>) -> Self {
64 Self(key.into())
65 }
66 pub fn as_str(&self) -> &str {
67 &self.0
68 }
69 pub fn is_empty(&self) -> bool {
70 self.0.is_empty()
71 }
72}
73
74impl std::fmt::Debug for CredentialRef {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 f.write_str("CredentialRef([REDACTED])")
77 }
78}
79
80#[derive(Debug, Clone)]
90pub struct ChatRequest {
91 pub model: String,
92 pub messages: Vec<ChatMessage>,
93 pub max_tokens: Option<u32>,
94 pub temperature_milli: u32,
96 pub credential: Option<CredentialRef>,
97}
98
99impl ChatRequest {
100 pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
101 Self {
102 model: model.into(),
103 messages,
104 max_tokens: None,
105 temperature_milli: 700,
106 credential: None,
107 }
108 }
109
110 pub fn with_max_tokens(mut self, n: u32) -> Self {
111 self.max_tokens = Some(n);
112 self
113 }
114
115 pub fn with_temperature(mut self, t_milli: u32) -> Self {
116 self.temperature_milli = t_milli;
117 self
118 }
119
120 pub fn with_credential(mut self, cred: CredentialRef) -> Self {
121 self.credential = Some(cred);
122 self
123 }
124}
125
126#[derive(Debug, Clone, PartialEq, Eq)]
128pub struct ChatToken {
129 pub text: String,
130 pub is_final: bool,
131}
132
133#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct ChatResponse {
136 pub content: String,
137 pub model: String,
138 pub prompt_tokens: u32,
139 pub completion_tokens: u32,
140}
141
142pub trait LlmProvider: Send + Sync {
146 fn chat(&self, req: &ChatRequest) -> crate::Result<ChatResponse>;
148
149 fn chat_stream(
153 &self,
154 req: &ChatRequest,
155 on_token: &mut dyn FnMut(ChatToken),
156 ) -> crate::Result<()>;
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162 use crate::Result;
163
164 struct EchoProvider;
165 impl LlmProvider for EchoProvider {
166 fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
167 let echo = req
168 .messages
169 .last()
170 .map(|m| m.content.as_str())
171 .unwrap_or("")
172 .to_owned();
173 Ok(ChatResponse {
174 content: echo.clone(),
175 model: req.model.clone(),
176 prompt_tokens: 1,
177 completion_tokens: 1,
178 })
179 }
180 fn chat_stream(
181 &self,
182 req: &ChatRequest,
183 on_token: &mut dyn FnMut(ChatToken),
184 ) -> Result<()> {
185 let text = req
186 .messages
187 .last()
188 .map(|m| m.content.as_str())
189 .unwrap_or("");
190 for ch in text.chars() {
191 on_token(ChatToken {
192 text: ch.to_string(),
193 is_final: false,
194 });
195 }
196 on_token(ChatToken {
197 text: String::new(),
198 is_final: true,
199 });
200 Ok(())
201 }
202 }
203
204 #[test]
205 fn chat_request_builder() {
206 let req = ChatRequest::new("local", vec![ChatMessage::user("hello")])
207 .with_max_tokens(256)
208 .with_temperature(500);
209 assert_eq!(req.max_tokens, Some(256));
210 assert_eq!(req.temperature_milli, 500);
211 }
212
213 #[test]
214 fn echo_provider_round_trips() {
215 let p = EchoProvider;
216 let req = ChatRequest::new("test", vec![ChatMessage::user("ping")]);
217 let resp = p.chat(&req).unwrap();
218 assert_eq!(resp.content, "ping");
219 }
220
221 #[test]
222 fn stream_delivers_all_chars_then_final() {
223 let p = EchoProvider;
224 let req = ChatRequest::new("test", vec![ChatMessage::user("hi")]);
225 let mut tokens: Vec<ChatToken> = Vec::new();
226 p.chat_stream(&req, &mut |t| tokens.push(t)).unwrap();
227 assert!(tokens.last().unwrap().is_final);
228 let text: String = tokens
229 .iter()
230 .filter(|t| !t.is_final)
231 .map(|t| t.text.as_str())
232 .collect();
233 assert_eq!(text, "hi");
234 }
235}