1use crate::models::Message;
2use serde::{Deserialize, Serialize};
3use std::{fmt::Display, time};
4
5#[derive(Default)]
6pub struct CodeContext {
7 pub language: String,
8 pub code: String,
9}
10
11#[derive(Debug, Clone)]
12pub struct Model {
13 id: String,
14 provider: String,
15}
16
17#[derive(Debug)]
18pub struct BackendResponse {
19 pub model: String,
20 pub id: String,
21 pub text: String,
22 pub done: bool,
23 pub init_conversation: bool,
24 pub usage: Option<BackendUsage>,
25}
26
27#[derive(Debug, Default, Clone)]
28pub struct BackendUsage {
29 pub prompt_tokens: usize,
30 pub completion_tokens: usize,
31 pub total_tokens: usize,
32}
33
34pub struct BackendPrompt {
35 model: String,
36 text: String,
37 context: Vec<Message>,
38 no_generate_title: bool,
39}
40
41impl BackendResponse {
42 pub fn new(id: &str, model: &str) -> Self {
43 Self {
44 model: model.into(),
45 id: id.into(),
46 text: String::new(),
47 done: false,
48 init_conversation: false,
49 usage: None,
50 }
51 }
52
53 pub fn with_text(mut self, text: impl Into<String>) -> Self {
54 self.text = text.into();
55 self
56 }
57
58 pub fn with_done(mut self) -> Self {
59 self.done = true;
60 self
61 }
62
63 pub fn with_init_conversation(mut self, init: bool) -> Self {
64 self.init_conversation = init;
65 self
66 }
67
68 pub fn with_usage(mut self, usage: BackendUsage) -> Self {
69 self.usage = Some(usage);
70 self
71 }
72}
73
74impl BackendPrompt {
75 pub fn new(text: impl Into<String>) -> BackendPrompt {
76 BackendPrompt {
77 model: String::new(),
78 text: text.into(),
79 context: vec![],
80 no_generate_title: false,
81 }
82 }
83
84 pub fn with_model(mut self, model: &str) -> Self {
85 self.model = model.to_string();
86 self
87 }
88
89 pub fn with_context(mut self, ctx: Vec<Message>) -> Self {
90 self.context = ctx;
91 self
92 }
93
94 pub fn with_no_generate_title(mut self) -> Self {
95 self.no_generate_title = true;
96 self
97 }
98
99 pub fn model(&self) -> &str {
100 &self.model
101 }
102
103 pub fn text(&self) -> &str {
104 &self.text
105 }
106
107 pub fn context(&self) -> &[Message] {
108 &self.context
109 }
110
111 pub fn no_generate_title(&self) -> bool {
112 self.no_generate_title
113 }
114}
115
116#[derive(Deserialize, Serialize, Debug, Clone)]
117pub struct BackendConnection {
118 #[serde(default)]
119 enabled: bool,
120 kind: BackendKind,
121 #[serde(default)]
122 alias: Option<String>,
123 #[serde(default)]
124 endpoint: String,
125 #[serde(default)]
126 api_key: Option<String>,
127 #[serde(default)]
128 timeout: Option<time::Duration>,
129 #[serde(default)]
130 models: Vec<String>,
131
132 #[serde(default)]
133 max_output_tokens: Option<usize>,
134}
135
136impl BackendConnection {
137 pub fn new(kind: BackendKind, endpoint: impl Into<String>) -> Self {
138 Self {
139 enabled: false,
140 kind,
141 alias: None,
142 endpoint: endpoint.into(),
143 api_key: None,
144 timeout: None,
145 models: Vec::new(),
146 max_output_tokens: None,
147 }
148 }
149
150 pub fn with_enabled(mut self, enabled: bool) -> Self {
151 self.enabled = enabled;
152 self
153 }
154
155 pub fn with_models(mut self, models: Vec<String>) -> Self {
156 self.models = models;
157 self
158 }
159
160 pub fn add_model(mut self, model: String) -> Self {
161 self.models.push(model);
162 self
163 }
164
165 pub fn with_alias(mut self, alias: impl Into<String>) -> Self {
166 self.alias = Some(alias.into());
167 self
168 }
169
170 pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
171 self.api_key = Some(api_key.into());
172 self
173 }
174
175 pub fn with_timeout(mut self, timeout: impl Into<time::Duration>) -> Self {
176 self.timeout = Some(timeout.into());
177 self
178 }
179
180 pub fn kind(&self) -> &BackendKind {
181 &self.kind
182 }
183
184 pub fn alias(&self) -> Option<&str> {
185 self.alias.as_deref()
186 }
187
188 pub fn endpoint(&self) -> &str {
189 &self.endpoint
190 }
191
192 pub fn api_key(&self) -> Option<&str> {
193 self.api_key.as_deref()
194 }
195
196 pub fn timeout(&self) -> Option<time::Duration> {
197 self.timeout
198 }
199
200 pub fn models(&self) -> &[String] {
201 &self.models
202 }
203
204 pub fn enabled(&self) -> bool {
205 self.enabled
206 }
207
208 pub fn max_output_tokens(&self) -> Option<usize> {
209 self.max_output_tokens
210 }
211}
212
213impl Model {
214 pub fn new(id: impl Into<String>) -> Self {
215 Self {
216 id: id.into(),
217 provider: String::new(),
218 }
219 }
220
221 pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
222 self.provider = provider.into();
223 self
224 }
225
226 pub fn id(&self) -> &str {
227 &self.id
228 }
229
230 pub fn provider(&self) -> &str {
231 &self.provider
232 }
233}
234
235#[derive(Hash, PartialEq, Eq, Deserialize, Serialize, Debug, Clone)]
236pub enum BackendKind {
237 #[serde(rename = "openai")]
238 OpenAI,
239 #[serde(rename = "gemini")]
240 Gemini,
241}
242
243impl Display for BackendKind {
244 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245 match self {
246 BackendKind::OpenAI => write!(f, "open_ai"),
247 BackendKind::Gemini => write!(f, "gemini"),
248 }
249 }
250}
251
252impl Display for BackendUsage {
253 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254 write!(
255 f,
256 "Prompt Tokens: {}, Completion Token: {}, Total: {}",
257 self.prompt_tokens, self.completion_tokens, self.total_tokens
258 )
259 }
260}
261
262impl Display for Model {
263 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
264 write!(
265 f,
266 "{} ({})",
267 self.id,
268 if self.provider.is_empty() {
269 "unknown"
270 } else {
271 &self.provider
272 }
273 )
274 }
275}