Skip to main content

aico/
models.rs

1use std::str::FromStr;
2
3use serde::{Deserialize, Serialize};
4use time::OffsetDateTime;
5
6use crate::exceptions::AicoError;
7
8// --- Enums ---
9
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum Mode {
13    Conversation,
14    Diff,
15    Raw,
16}
17
18impl std::fmt::Display for Mode {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        match self {
21            Mode::Conversation => write!(f, "conversation"),
22            Mode::Diff => write!(f, "diff"),
23            Mode::Raw => write!(f, "raw"),
24        }
25    }
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29#[serde(rename_all = "lowercase")]
30pub enum Role {
31    User,
32    Assistant,
33    System,
34}
35
36impl std::fmt::Display for Role {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        match self {
39            Role::User => write!(f, "user"),
40            Role::Assistant => write!(f, "assistant"),
41            Role::System => write!(f, "system"),
42        }
43    }
44}
45
46#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
47#[serde(rename_all = "lowercase")]
48pub enum Provider {
49    OpenAI,
50    OpenRouter,
51}
52
53impl FromStr for Provider {
54    type Err = AicoError;
55    fn from_str(s: &str) -> Result<Self, Self::Err> {
56        match s {
57            "openai" => Ok(Self::OpenAI),
58            "openrouter" => Ok(Self::OpenRouter),
59            _ => Err(AicoError::Configuration(format!(
60                "Unrecognized provider prefix in '{}'. Use 'openai/' or 'openrouter/'.",
61                s
62            ))),
63        }
64    }
65}
66
67// --- Shared History Models (historystore/models.py) ---
68
69#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
70pub struct HistoryRecord {
71    pub role: Role,
72    pub content: String,
73    pub mode: Mode,
74    #[serde(with = "time::serde::rfc3339", default = "default_timestamp")]
75    pub timestamp: OffsetDateTime,
76
77    #[serde(default)]
78    pub passthrough: bool,
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub piped_content: Option<String>,
81
82    // Assistant-only optional metadata
83    #[serde(default, skip_serializing_if = "Option::is_none")]
84    pub model: Option<String>,
85    #[serde(default, skip_serializing_if = "Option::is_none")]
86    pub token_usage: Option<TokenUsage>,
87    #[serde(default, skip_serializing_if = "Option::is_none")]
88    pub cost: Option<f64>,
89    #[serde(default, skip_serializing_if = "Option::is_none")]
90    pub duration_ms: Option<u64>,
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub derived: Option<DerivedContent>,
93
94    // Edit lineage
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub edit_of: Option<usize>,
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
100pub struct SessionView {
101    pub model: String,
102    #[serde(default)]
103    pub context_files: Vec<String>,
104    #[serde(default)]
105    pub message_indices: Vec<usize>,
106    #[serde(default)]
107    pub history_start_pair: usize,
108    #[serde(default)]
109    pub excluded_pairs: Vec<usize>,
110    #[serde(with = "time::serde::rfc3339", default = "default_timestamp")]
111    pub created_at: OffsetDateTime,
112}
113
114// --- Session Pointer (.ai_session.json) ---
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct SessionPointer {
118    #[serde(rename = "type")]
119    pub pointer_type: String, // "aico_session_pointer_v1"
120    pub path: String,
121}
122
123// --- Supporting Structs ---
124
125#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
126pub struct TokenUsage {
127    pub prompt_tokens: u32,
128    pub completion_tokens: u32,
129    pub total_tokens: u32,
130    #[serde(default, skip_serializing_if = "Option::is_none")]
131    pub cached_tokens: Option<u32>,
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub reasoning_tokens: Option<u32>,
134    #[serde(default, skip_serializing_if = "Option::is_none")]
135    pub cost: Option<f64>,
136}
137
138#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
139pub struct DerivedContent {
140    #[serde(default, skip_serializing_if = "Option::is_none")]
141    pub unified_diff: Option<String>,
142    #[serde(
143        default,
144        skip_serializing_if = "Vec::is_empty",
145        deserialize_with = "deserialize_null_as_default"
146    )]
147    pub display_content: Vec<DisplayItem>,
148}
149
150fn deserialize_null_as_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
151where
152    D: serde::Deserializer<'de>,
153    T: Default + serde::Deserialize<'de>,
154{
155    let opt = Option::deserialize(deserializer)?;
156    Ok(opt.unwrap_or_default())
157}
158
159#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
160#[serde(tag = "type", content = "content", rename_all = "lowercase")]
161pub enum DisplayItem {
162    #[serde(alias = "text")]
163    Markdown(String),
164    Diff(String),
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct MessageWithId {
169    #[serde(flatten)]
170    pub record: HistoryRecord,
171    pub id: usize,
172}
173
174#[derive(Debug, Clone)]
175pub struct MessageWithContext {
176    pub record: HistoryRecord,
177    pub global_index: usize,
178    pub pair_index: usize,
179    pub is_excluded: bool,
180}
181
182#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct MessagePairJson {
184    pub pair_index: usize,
185    pub user: MessageWithId,
186    pub assistant: MessageWithId,
187}
188
189#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
190pub enum AddonSource {
191    Project,
192    User,
193    Bundled,
194}
195
196impl std::fmt::Display for AddonSource {
197    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        match self {
199            AddonSource::Project => write!(f, "project"),
200            AddonSource::User => write!(f, "user"),
201            AddonSource::Bundled => write!(f, "bundled"),
202        }
203    }
204}
205
206#[derive(Debug, Clone)]
207pub struct AddonInfo {
208    pub name: String,
209    pub path: std::path::PathBuf,
210    pub help_text: String,
211    pub source: AddonSource,
212}
213
214#[derive(Debug, Clone, Serialize, Deserialize)]
215pub struct TokenInfo {
216    pub description: String,
217    pub tokens: u32,
218    pub cost: Option<f64>,
219}
220
221pub struct ActiveWindowSummary {
222    pub active_pairs: usize,
223    pub active_start_id: usize,
224    pub active_end_id: usize,
225    pub excluded_in_window: usize,
226    pub pairs_sent: usize,
227    pub has_dangling: bool,
228}
229
230// --- Diffing / Streaming Models ---
231
232#[derive(Debug, Clone, PartialEq)]
233pub struct AIPatch {
234    pub llm_file_path: String,
235    pub search_content: String,
236    pub replace_content: String,
237    pub indent: String,
238    pub raw_block: String,
239}
240
241#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
242pub struct ProcessedDiffBlock {
243    pub llm_file_path: String,
244    pub unified_diff: String,
245}
246
247#[derive(Debug, Clone, PartialEq)]
248pub struct FileHeader {
249    pub llm_file_path: String,
250}
251
252#[derive(Debug, Clone, PartialEq)]
253pub struct WarningMessage {
254    pub text: String,
255}
256
257#[derive(Debug, Clone, PartialEq)]
258pub struct UnparsedBlock {
259    pub text: String,
260}
261
262#[derive(Debug, Clone, PartialEq)]
263pub enum StreamYieldItem {
264    Text(String),
265    IncompleteBlock(String),
266    FileHeader(FileHeader),
267    DiffBlock(ProcessedDiffBlock),
268    Patch(AIPatch),
269    Warning(WarningMessage),
270    Unparsed(UnparsedBlock),
271}
272
273impl StreamYieldItem {
274    pub fn is_warning(&self) -> bool {
275        matches!(self, StreamYieldItem::Warning(_))
276    }
277
278    pub fn to_display_item(self, is_final: bool) -> Option<DisplayItem> {
279        match self {
280            StreamYieldItem::Text(t) => Some(DisplayItem::Markdown(t)),
281            StreamYieldItem::FileHeader(h) => Some(DisplayItem::Markdown(format!(
282                "File: `{}`\n",
283                h.llm_file_path
284            ))),
285            StreamYieldItem::DiffBlock(db) => Some(DisplayItem::Diff(db.unified_diff)),
286            StreamYieldItem::Warning(w) => {
287                Some(DisplayItem::Markdown(format!("[!WARNING]\n{}\n\n", w.text)))
288            }
289            StreamYieldItem::Unparsed(u) => Some(DisplayItem::Markdown(format!(
290                "\n`````text\n{}\n`````\n",
291                u.text
292            ))),
293            StreamYieldItem::IncompleteBlock(t) => {
294                if is_final {
295                    Some(DisplayItem::Markdown(t))
296                } else {
297                    None
298                }
299            }
300            StreamYieldItem::Patch(_) => None,
301        }
302    }
303}
304
305#[derive(Debug, Serialize, Deserialize)]
306pub struct StatusResponse {
307    pub session_name: String,
308    pub model: String,
309    pub context_files: Vec<String>,
310    #[serde(skip_serializing_if = "Option::is_none")]
311    pub total_tokens: Option<u32>,
312    #[serde(skip_serializing_if = "Option::is_none")]
313    pub total_cost: Option<f64>,
314}
315
316#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
317pub struct InteractionResult {
318    pub content: String,
319    pub display_items: Option<Vec<DisplayItem>>,
320    pub token_usage: Option<TokenUsage>,
321    pub cost: Option<f64>,
322    pub duration_ms: u64,
323    pub unified_diff: Option<String>,
324}
325
326#[derive(Debug, Clone)]
327pub struct InteractionConfig {
328    pub mode: Mode,
329    pub no_history: bool,
330    pub passthrough: bool,
331    pub model_override: Option<String>,
332}
333
334#[derive(Debug, Clone)]
335pub struct ContextState<'a> {
336    pub static_files: Vec<(&'a str, &'a str)>,
337    pub floating_files: Vec<(&'a str, &'a str)>,
338    pub splice_idx: usize,
339}
340
341// --- Helpers ---
342
343pub fn default_timestamp() -> OffsetDateTime {
344    OffsetDateTime::now_utc()
345}