mod.rs

crates/cortex-core/src/storage/mod.rs

1. //! Storage: one SQLite file with full-text search over messages and titles.
3. use std::path::Path;
5. use chrono::{Datelike, TimeZone};
6. use rusqlite::{params, Connection, OptionalExtension, Row};
7. use serde::{Deserialize, Serialize};
8. use ts_rs::TS;
10. use crate::pricing::ModelInfo;
11. use crate::types::{Prices, Usage, UsageSource};
13. mod search;
14. pub use search::fts_query;
16. const MIGRATIONS: &[&str] = &[include_str!("migrations/0001_init.sql"), include_str!("migrations/0002_fts.sql")];
18. #[derive(Debug, thiserror::Error)]
19. pub enum DbError {
20. #[error(transparent)]
21. Sqlite(#[from] rusqlite::Error),
22. #[error(transparent)]
23. Json(#[from] serde_json::Error),
24. #[error("not found: {0}")]
25. NotFound(String),
26. }
28. pub type Result<T> = std::result::Result<T, DbError>;
30. pub fn now_ms() -> i64 {
31. chrono::Utc::now().timestamp_millis()
32. }
34. pub fn new_id() -> String {
35. uuid::Uuid::now_v7().to_string()
36. }
38. /// Start of the calendar month containing `now`, in `now`'s own time zone, as UTC ms.
39. pub fn month_start_ms<Tz: TimeZone>(now: &chrono::DateTime<Tz>) -> i64 {
40. let tz = now.timezone();
41. tz.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
42. .earliest()
43. .map(|d| d.timestamp_millis())
44. .unwrap_or_else(|| now.timestamp_millis())
45. }
47. // ---- types shared with the frontend ---------------------------------------
49. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
50. #[ts(export)]
51. pub struct Provider {
52. pub id: String,
53. /// `deepinfra` or `openai_compatible`
54. pub kind: String,
55. pub name: String,
56. pub base_url: String,
57. pub requires_key: bool,
58. }
60. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
61. #[ts(export)]
62. pub struct Conversation {
63. pub id: String,
64. pub title: String,
65. pub provider_id: String,
66. pub model_id: String,
67. pub preset_id: Option<String>,
68. pub system_prompt: String,
69. pub origin: String,
70. pub pinned: bool,
71. pub archived: bool,
72. #[ts(type = "number")]
73. pub created_at: i64,
74. #[ts(type = "number")]
75. pub updated_at: i64,
76. }
78. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
79. #[ts(export)]
80. pub struct ConversationSummary {
81. pub id: String,
82. pub title: String,
83. pub model_id: String,
84. pub pinned: bool,
85. pub origin: String,
86. #[ts(type = "number")]
87. pub updated_at: i64,
88. #[ts(type = "number")]
89. pub cost_nanousd: i64,
90. pub cost_estimated: bool,
91. #[ts(type = "number")]
92. pub message_count: i64,
93. /// Some reply was answered (tokens counted), so $0 means free, not failed.
94. pub has_usage: bool,
95. }
97. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
98. #[ts(export)]
99. pub struct AttachmentMeta {
100. pub id: String,
101. /// `image` or `text`
102. pub kind: String,
103. pub filename: String,
104. pub mime: String,
105. #[ts(type = "number")]
106. pub byte_size: i64,
107. pub width: Option<u32>,
108. pub height: Option<u32>,
109. }
111. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
112. #[ts(export)]
113. pub struct Message {
114. pub id: String,
115. pub conversation_id: String,
116. pub role: String,
117. pub content: String,
118. pub reasoning: Option<String>,
119. /// `complete`, `streaming`, `cancelled` or `error`
120. pub status: String,
121. pub error_kind: Option<String>,
122. pub error_message: Option<String>,
123. pub provider_id: Option<String>,
124. pub model_id: Option<String>,
125. pub usage: Usage,
126. pub usage_source: UsageSource,
127. pub prices: Option<Prices>,
128. #[ts(type = "number")]
129. pub cost_nanousd: i64,
130. pub finish_reason: Option<String>,
131. #[ts(type = "number | null")]
132. pub ttft_ms: Option<i64>,
133. #[ts(type = "number | null")]
134. pub duration_ms: Option<i64>,
135. #[ts(type = "number")]
136. pub created_at: i64,
137. pub attachments: Vec<AttachmentMeta>,
138. }
140. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
141. #[ts(export)]
142. pub struct Preset {
143. pub id: String,
144. pub name: String,
145. pub system_prompt: String,
146. pub provider_id: Option<String>,
147. pub model_id: Option<String>,
148. pub temperature: Option<f32>,
149. pub max_tokens: Option<u32>,
150. }
152. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
153. #[ts(export)]
154. pub struct SearchHit {
155. pub conversation_id: String,
156. pub title: String,
157. pub message_id: Option<String>,
158. /// With `\u{1}` and `\u{2}` around matched terms, for the UI to highlight.
159. pub snippet: String,
160. #[ts(type = "number")]
161. pub created_at: i64,
162. }
164. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
165. #[ts(export)]
166. pub struct UsageRow {
167. /// A model id, or a day as YYYY-MM-DD.
168. pub key: String,
169. #[ts(type = "number")]
170. pub cost_nanousd: i64,
171. #[ts(type = "number")]
172. pub prompt_tokens: i64,
173. #[ts(type = "number")]
174. pub completion_tokens: i64,
175. #[ts(type = "number")]
176. pub messages: i64,
177. pub estimated: bool,
178. }
180. /// What the stream produced, written onto the assistant row in one go.
181. #[derive(Debug, Clone, Default)]
182. pub struct Finalize {
183. pub content: String,
184. pub reasoning: Option<String>,
185. pub status: &'static str,
186. pub error_kind: Option<String>,
187. pub error_message: Option<String>,
188. pub usage: Usage,
189. pub usage_source: Option<UsageSource>,
190. pub prices: Option<Prices>,
191. pub cost_nanousd: i64,
192. pub provider_cost_nanousd: Option<i64>,
193. pub finish_reason: Option<String>,
194. pub ttft_ms: Option<i64>,
195. pub duration_ms: Option<i64>,
196. }
198. pub struct NewAttachment<'a> {
199. pub kind: &'a str,
200. pub filename: &'a str,
201. pub mime: &'a str,
202. pub byte_size: i64,
203. pub sha256: &'a str,
204. pub width: Option<u32>,
205. pub height: Option<u32>,
206. pub path: Option<&'a str>,
207. pub text_content: Option<&'a str>,
208. }
210. /// An attachment as the provider request needs it.
211. #[derive(Debug, Clone)]
212. pub struct AttachmentData {
213. pub meta: AttachmentMeta,
214. pub path: Option<String>,
215. pub text_content: Option<String>,
216. }
218. // ---- the database ----------------------------------------------------------
220. pub struct Db {
221. conn: Connection,
222. }
224. const MESSAGE_COLS: &str = "id, conversation_id, role, content, reasoning, status, error_kind, error_message, \
225. provider_id, model_id, prompt_tokens, completion_tokens, cached_tokens, usage_source, \
226. price_in_per_m, price_out_per_m, price_cache_per_m, cost_nanousd, finish_reason, ttft_ms, duration_ms, created_at";
228. const CONVERSATION_COLS: &str =
229. "id, title, provider_id, model_id, preset_id, system_prompt, origin, pinned, archived, created_at, updated_at";
231. fn conversation_row(r: &Row) -> rusqlite::Result<Conversation> {
232. Ok(Conversation {
233. id: r.get(0)?,
234. title: r.get(1)?,
235. provider_id: r.get(2)?,
236. model_id: r.get(3)?,
237. preset_id: r.get(4)?,
238. system_prompt: r.get(5)?,
239. origin: r.get(6)?,
240. pinned: r.get(7)?,
241. archived: r.get(8)?,
242. created_at: r.get(9)?,
243. updated_at: r.get(10)?,
244. })
245. }
247. fn message_row(r: &Row) -> rusqlite::Result<Message> {
248. let pin: Option<f64> = r.get(14)?;
249. let pout: Option<f64> = r.get(15)?;
250. let pcache: Option<f64> = r.get(16)?;
251. Ok(Message {
252. id: r.get(0)?,
253. conversation_id: r.get(1)?,
254. role: r.get(2)?,
255. content: r.get(3)?,
256. reasoning: r.get(4)?,
257. status: r.get(5)?,
258. error_kind: r.get(6)?,
259. error_message: r.get(7)?,
260. provider_id: r.get(8)?,
261. model_id: r.get(9)?,
262. usage: Usage { prompt_tokens: r.get(10)?, completion_tokens: r.get(11)?, cached_tokens: r.get(12)? },
263. usage_source: UsageSource::parse(&r.get::<_, String>(13)?),
264. prices: match (pin, pout) {
265. (Some(i), Some(o)) => Some(Prices { input_per_m: i, output_per_m: o, cache_read_per_m: pcache }),
266. _ => None,
267. },
268. cost_nanousd: r.get(17)?,
269. finish_reason: r.get(18)?,
270. ttft_ms: r.get(19)?,
271. duration_ms: r.get(20)?,
272. created_at: r.get(21)?,
273. attachments: vec![],
274. })
275. }
277. fn attachment_meta_row(r: &Row) -> rusqlite::Result<AttachmentMeta> {
278. Ok(AttachmentMeta {
279. id: r.get(0)?,
280. kind: r.get(1)?,
281. filename: r.get(2)?,
282. mime: r.get(3)?,
283. byte_size: r.get(4)?,
284. width: r.get(5)?,
285. height: r.get(6)?,
286. })
287. }
289. pub const DEEPINFRA_ID: &str = "deepinfra";
291. impl Db {
292. pub fn open(path: &Path) -> Result<Self> {
293. Self::init(Connection::open(path)?)
294. }
296. pub fn open_in_memory() -> Result<Self> {
297. Self::init(Connection::open_in_memory()?)
298. }
300. fn init(conn: Connection) -> Result<Self> {
301. conn.pragma_update(None, "journal_mode", "WAL")?;
302. conn.pragma_update(None, "foreign_keys", "ON")?;
303. conn.busy_timeout(std::time::Duration::from_secs(5))?;
304. let mut db = Self { conn };
305. db.migrate()?;
306. db.seed()?;
307. db.recover_interrupted()?;
308. Ok(db)
309. }
311. fn migrate(&mut self) -> Result<()> {
312. let version: i64 = self.conn.pragma_query_value(None, "user_version", |r| r.get(0))?;
313. for (i, sql) in MIGRATIONS.iter().enumerate().skip(version as usize) {
314. let tx = self.conn.transaction()?;
315. tx.execute_batch(sql)?;
316. tx.pragma_update(None, "user_version", (i + 1) as i64)?;
317. tx.commit()?;
318. }
319. Ok(())
320. }
322. #[cfg(test)]
323. pub(crate) fn schema_version(&self) -> Result<i64> {
324. Ok(self.conn.pragma_query_value(None, "user_version", |r| r.get(0))?)
325. }
327. fn seed(&self) -> Result<()> {
328. self.conn.execute(
329. "INSERT OR IGNORE INTO providers (id, kind, name, base_url, requires_key, created_at)
330. VALUES (?1, 'deepinfra', 'DeepInfra', ?2, 1, ?3)",
331. params![DEEPINFRA_ID, crate::provider::DEEPINFRA_BASE, now_ms()],
332. )?;
333. Ok(())
334. }
336. /// A crash or quit mid-stream leaves a row marked streaming; close it off.
337. fn recover_interrupted(&self) -> Result<usize> {
338. Ok(self.conn.execute(
339. "UPDATE messages SET status = 'error', error_kind = 'interrupted',
340. error_message = 'Cortex closed before this reply finished.'
341. WHERE status = 'streaming'",
342. [],
343. )?)
344. }
346. // ---- providers ----
348. pub fn list_providers(&self) -> Result<Vec<Provider>> {
349. let mut st = self.conn.prepare(
350. "SELECT id, kind, name, base_url, requires_key FROM providers ORDER BY kind = 'deepinfra' DESC, created_at",
351. )?;
352. let rows = st.query_map([], |r| {
353. Ok(Provider {
354. id: r.get(0)?,
355. kind: r.get(1)?,
356. name: r.get(2)?,
357. base_url: r.get(3)?,
358. requires_key: r.get(4)?,
359. })
360. })?;
361. Ok(rows.collect::<rusqlite::Result<_>>()?)
362. }
364. pub fn get_provider(&self, id: &str) -> Result<Provider> {
365. self.list_providers()?
366. .into_iter()
367. .find(|p| p.id == id)
368. .ok_or_else(|| DbError::NotFound(format!("provider {id}")))
369. }
371. pub fn upsert_provider(&self, p: &Provider) -> Result<()> {
372. self.conn.execute(
373. "INSERT INTO providers (id, kind, name, base_url, requires_key, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
374. ON CONFLICT (id) DO UPDATE SET name = excluded.name, base_url = excluded.base_url,
375. requires_key = excluded.requires_key",
376. params![p.id, p.kind, p.name, p.base_url.trim_end_matches('/'), p.requires_key, now_ms()],
377. )?;
378. Ok(())
379. }
381. pub fn delete_provider(&self, id: &str) -> Result<()> {
382. if id == DEEPINFRA_ID {
383. return Ok(()); // the built-in provider stays
384. }
385. self.conn.execute("DELETE FROM providers WHERE id = ?1", [id])?;
386. self.conn.execute("DELETE FROM model_pricing WHERE provider_id = ?1", [id])?;
387. Ok(())
388. }
390. // ---- conversations ----
392. pub fn create_conversation(
393. &self,
394. provider_id: &str,
395. model_id: &str,
396. system_prompt: &str,
397. preset_id: Option<&str>,
398. origin: &str,
399. ) -> Result<Conversation> {
400. let id = new_id();
401. let now = now_ms();
402. self.conn.execute(
403. "INSERT INTO conversations (id, title, provider_id, model_id, preset_id, system_prompt, origin, created_at, updated_at)
404. VALUES (?1, 'New chat', ?2, ?3, ?4, ?5, ?6, ?7, ?7)",
405. params![id, provider_id, model_id, preset_id, system_prompt, origin, now],
406. )?;
407. self.get_conversation(&id)
408. }
410. pub fn get_conversation(&self, id: &str) -> Result<Conversation> {
411. self.conn
412. .query_row(&format!("SELECT {CONVERSATION_COLS} FROM conversations WHERE id = ?1"), [id], conversation_row)
413. .optional()?
414. .ok_or_else(|| DbError::NotFound(format!("conversation {id}")))
415. }
417. pub fn list_conversations(&self, include_archived: bool, limit: i64) -> Result<Vec<ConversationSummary>> {
418. let mut st = self.conn.prepare(
419. "SELECT c.id, c.title, c.model_id, c.pinned, c.origin, c.updated_at,
420. COALESCE(SUM(m.cost_nanousd), 0),
421. COALESCE(MAX(m.usage_source = 'estimated'), 0),
422. COUNT(m.id),
423. COALESCE(MAX(m.role = 'assistant' AND m.usage_source != 'none'), 0)
424. FROM conversations c LEFT JOIN messages m ON m.conversation_id = c.id
425. WHERE (?1 OR c.archived = 0)
426. GROUP BY c.id ORDER BY c.pinned DESC, c.updated_at DESC LIMIT ?2",
427. )?;
428. let rows = st.query_map(params![include_archived, limit], |r| {
429. Ok(ConversationSummary {
430. id: r.get(0)?,
431. title: r.get(1)?,
432. model_id: r.get(2)?,
433. pinned: r.get(3)?,
434. origin: r.get(4)?,
435. updated_at: r.get(5)?,
436. cost_nanousd: r.get(6)?,
437. cost_estimated: r.get(7)?,
438. message_count: r.get(8)?,
439. has_usage: r.get(9)?,
440. })
441. })?;
442. Ok(rows.collect::<rusqlite::Result<_>>()?)
443. }
445. pub fn rename_conversation(&self, id: &str, title: &str) -> Result<()> {
446. let title = title.trim();
447. let title = if title.is_empty() { "New chat" } else { title };
448. self.conn.execute("UPDATE conversations SET title = ?2 WHERE id = ?1", params![id, title])?;
449. Ok(())
450. }
452. pub fn set_pinned(&self, id: &str, pinned: bool) -> Result<()> {
453. self.conn.execute("UPDATE conversations SET pinned = ?2 WHERE id = ?1", params![id, pinned])?;
454. Ok(())
455. }
457. pub fn set_archived(&self, id: &str, archived: bool) -> Result<()> {
458. self.conn.execute("UPDATE conversations SET archived = ?2 WHERE id = ?1", params![id, archived])?;
459. Ok(())
460. }
462. pub fn set_conversation_model(&self, id: &str, provider_id: &str, model_id: &str) -> Result<()> {
463. self.conn.execute(
464. "UPDATE conversations SET provider_id = ?2, model_id = ?3 WHERE id = ?1",
465. params![id, provider_id, model_id],
466. )?;
467. Ok(())
468. }
470. pub fn set_system_prompt(&self, id: &str, system_prompt: &str, preset_id: Option<&str>) -> Result<()> {
471. self.conn.execute(
472. "UPDATE conversations SET system_prompt = ?2, preset_id = ?3 WHERE id = ?1",
473. params![id, system_prompt, preset_id],
474. )?;
475. Ok(())
476. }
478. pub fn delete_conversation(&self, id: &str) -> Result<()> {
479. self.conn.execute("DELETE FROM conversations WHERE id = ?1", [id])?;
480. Ok(())
481. }
483. fn touch(&self, conversation_id: &str) -> Result<()> {
484. self.conn
485. .execute("UPDATE conversations SET updated_at = ?2 WHERE id = ?1", params![conversation_id, now_ms()])?;
486. Ok(())
487. }
489. // ---- messages ----
491. pub fn insert_user_message(
492. &self,
493. conversation_id: &str,
494. content: &str,
495. attachment_ids: &[String],
496. ) -> Result<Message> {
497. let id = new_id();
498. let tx = self.conn.unchecked_transaction()?;
499. tx.execute(
500. "INSERT INTO messages (id, conversation_id, role, content, status, created_at) VALUES (?1, ?2, 'user', ?3, 'complete', ?4)",
501. params![id, conversation_id, content, now_ms()],
502. )?;
503. for a in attachment_ids {
504. tx.execute("UPDATE attachments SET message_id = ?1 WHERE id = ?2 AND message_id IS NULL", params![id, a])?;
505. }
506. // The first message names the chat; no model call is spent on a title.
507. tx.execute(
508. "UPDATE conversations SET title = ?2 WHERE id = ?1 AND title = 'New chat'",
509. params![conversation_id, title_from(content)],
510. )?;
511. tx.commit()?;
512. self.touch(conversation_id)?;
513. self.get_message(&id)
514. }
516. pub fn insert_assistant_placeholder(
517. &self,
518. conversation_id: &str,
519. provider_id: &str,
520. model_id: &str,
521. ) -> Result<Message> {
522. let id = new_id();
523. self.conn.execute(
524. "INSERT INTO messages (id, conversation_id, role, content, status, provider_id, model_id, created_at)
525. VALUES (?1, ?2, 'assistant', '', 'streaming', ?3, ?4, ?5)",
526. params![id, conversation_id, provider_id, model_id, now_ms()],
527. )?;
528. self.get_message(&id)
529. }
531. pub fn finalize_assistant(&self, id: &str, f: &Finalize) -> Result<Message> {
532. let (pin, pout, pcache) = match f.prices {
533. Some(p) => (Some(p.input_per_m), Some(p.output_per_m), p.cache_read_per_m),
534. None => (None, None, None),
535. };
536. self.conn.execute(
537. "UPDATE messages SET content = ?2, reasoning = ?3, status = ?4, error_kind = ?5, error_message = ?6,
538. prompt_tokens = ?7, completion_tokens = ?8, cached_tokens = ?9, usage_source = ?10,
539. price_in_per_m = ?11, price_out_per_m = ?12, price_cache_per_m = ?13, cost_nanousd = ?14,
540. provider_cost_nanousd = ?15, finish_reason = ?16, ttft_ms = ?17, duration_ms = ?18
541. WHERE id = ?1",
542. params![
543. id,
544. f.content,
545. f.reasoning.as_deref().filter(|s| !s.is_empty()),
546. f.status,
547. f.error_kind,
548. f.error_message,
549. f.usage.prompt_tokens,
550. f.usage.completion_tokens,
551. f.usage.cached_tokens,
552. f.usage_source.unwrap_or(UsageSource::None).as_str(),
553. pin,
554. pout,
555. pcache,
556. f.cost_nanousd,
557. f.provider_cost_nanousd,
558. f.finish_reason,
559. f.ttft_ms,
560. f.duration_ms
561. ],
562. )?;
563. let m = self.get_message(id)?;
564. self.touch(&m.conversation_id)?;
565. Ok(m)
566. }
568. pub fn get_message(&self, id: &str) -> Result<Message> {
569. let mut m = self
570. .conn
571. .query_row(&format!("SELECT {MESSAGE_COLS} FROM messages WHERE id = ?1"), [id], message_row)
572. .optional()?
573. .ok_or_else(|| DbError::NotFound(format!("message {id}")))?;
574. m.attachments = self.attachments_for(&m.id)?;
575. Ok(m)
576. }
578. pub fn list_messages(&self, conversation_id: &str) -> Result<Vec<Message>> {
579. let mut st = self.conn.prepare(&format!(
580. "SELECT {MESSAGE_COLS} FROM messages WHERE conversation_id = ?1 ORDER BY created_at, rowid"
581. ))?;
582. let mut msgs: Vec<Message> = st.query_map([conversation_id], message_row)?.collect::<rusqlite::Result<_>>()?;
583. for m in &mut msgs {
584. m.attachments = self.attachments_for(&m.id)?;
585. }
586. Ok(msgs)
587. }
589. /// Delete a message and everything after it (edit and regenerate start here).
590. /// The message's own attachments go back to being drafts, so an edited
591. /// message can be sent again with them.
592. pub fn truncate_from(&self, conversation_id: &str, message_id: &str) -> Result<usize> {
593. let (created, rowid): (i64, i64) = self.conn.query_row(
594. "SELECT created_at, rowid FROM messages WHERE id = ?1 AND conversation_id = ?2",
595. params![message_id, conversation_id],
596. |r| Ok((r.get(0)?, r.get(1)?)),
597. )?;
598. self.conn.execute(
599. "UPDATE attachments SET message_id = NULL, created_at = ?2 WHERE message_id = ?1",
600. params![message_id, now_ms()],
601. )?;
602. Ok(self.conn.execute(
603. "DELETE FROM messages WHERE conversation_id = ?1 AND (created_at > ?2 OR (created_at = ?2 AND rowid >= ?3))",
604. params![conversation_id, created, rowid],
605. )?)
606. }
608. // ---- attachments ----
610. pub fn insert_attachment(&self, a: &NewAttachment) -> Result<AttachmentMeta> {
611. let id = new_id();
612. self.conn.execute(
613. "INSERT INTO attachments (id, kind, filename, mime, byte_size, sha256, width, height, path, text_content, created_at)
614. VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
615. params![id, a.kind, a.filename, a.mime, a.byte_size, a.sha256, a.width, a.height, a.path, a.text_content, now_ms()],
616. )?;
617. Ok(self.attachment(&id)?.meta)
618. }
620. pub fn attachment(&self, id: &str) -> Result<AttachmentData> {
621. self.conn
622. .query_row(
623. "SELECT id, kind, filename, mime, byte_size, width, height, path, text_content FROM attachments WHERE id = ?1",
624. [id],
625. |r| Ok(AttachmentData { meta: attachment_meta_row(r)?, path: r.get(7)?, text_content: r.get(8)? }),
626. )
627. .optional()?
628. .ok_or_else(|| DbError::NotFound(format!("attachment {id}")))
629. }
631. pub fn attachments_for(&self, message_id: &str) -> Result<Vec<AttachmentMeta>> {
632. let mut st = self.conn.prepare(
633. "SELECT id, kind, filename, mime, byte_size, width, height FROM attachments WHERE message_id = ?1 ORDER BY created_at, rowid",
634. )?;
635. let rows = st.query_map([message_id], attachment_meta_row)?;
636. Ok(rows.collect::<rusqlite::Result<_>>()?)
637. }
639. pub fn attachment_data_for(&self, message_id: &str) -> Result<Vec<AttachmentData>> {
640. self.attachments_for(message_id)?.into_iter().map(|a| self.attachment(&a.id)).collect()
641. }
643. /// Paths still referenced by any attachment row (a file may be shared by hash).
644. pub fn attachment_path_in_use(&self, path: &str) -> Result<bool> {
645. Ok(self.conn.query_row("SELECT EXISTS (SELECT 1 FROM attachments WHERE path = ?1)", [path], |r| r.get(0))?)
646. }
648. /// Remove an attachment that was never sent; returns its file path, if any.
649. pub fn delete_draft_attachment(&self, id: &str) -> Result<Option<String>> {
650. let path: Option<Option<String>> = self
651. .conn
652. .query_row("SELECT path FROM attachments WHERE id = ?1 AND message_id IS NULL", [id], |r| r.get(0))
653. .optional()?;
654. self.conn.execute("DELETE FROM attachments WHERE id = ?1 AND message_id IS NULL", [id])?;
655. Ok(path.flatten())
656. }
658. /// Drafts never sent (the composer was cleared or the app closed).
659. pub fn delete_stale_drafts(&self, older_than_ms: i64) -> Result<Vec<String>> {
660. let cutoff = now_ms() - older_than_ms;
661. let mut st = self.conn.prepare(
662. "SELECT path FROM attachments WHERE message_id IS NULL AND created_at < ?1 AND path IS NOT NULL",
663. )?;
664. let paths: Vec<String> = st.query_map([cutoff], |r| r.get(0))?.collect::<rusqlite::Result<_>>()?;
665. self.conn.execute("DELETE FROM attachments WHERE message_id IS NULL AND created_at < ?1", [cutoff])?;
666. Ok(paths)
667. }
669. // ---- presets ----
671. pub fn list_presets(&self) -> Result<Vec<Preset>> {
672. let mut st = self.conn.prepare(
673. "SELECT id, name, system_prompt, provider_id, model_id, temperature, max_tokens FROM presets ORDER BY sort_order, name",
674. )?;
675. let rows = st.query_map([], |r| {
676. Ok(Preset {
677. id: r.get(0)?,
678. name: r.get(1)?,
679. system_prompt: r.get(2)?,
680. provider_id: r.get(3)?,
681. model_id: r.get(4)?,
682. temperature: r.get(5)?,
683. max_tokens: r.get(6)?,
684. })
685. })?;
686. Ok(rows.collect::<rusqlite::Result<_>>()?)
687. }
689. pub fn save_preset(&self, p: &Preset) -> Result<Preset> {
690. let id = if p.id.is_empty() { new_id() } else { p.id.clone() };
691. let now = now_ms();
692. self.conn.execute(
693. "INSERT INTO presets (id, name, system_prompt, provider_id, model_id, temperature, max_tokens, created_at, updated_at)
694. VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8)
695. ON CONFLICT (id) DO UPDATE SET name = excluded.name, system_prompt = excluded.system_prompt,
696. provider_id = excluded.provider_id, model_id = excluded.model_id, temperature = excluded.temperature,
697. max_tokens = excluded.max_tokens, updated_at = excluded.updated_at",
698. params![id, p.name.trim(), p.system_prompt, p.provider_id, p.model_id, p.temperature, p.max_tokens, now],
699. )?;
700. Ok(Preset { id, ..p.clone() })
701. }
703. pub fn delete_preset(&self, id: &str) -> Result<()> {
704. self.conn.execute("DELETE FROM presets WHERE id = ?1", [id])?;
705. Ok(())
706. }
708. // ---- settings ----
710. pub fn get_setting<T: serde::de::DeserializeOwned>(&self, key: &str) -> Result<Option<T>> {
711. let raw: Option<String> =
712. self.conn.query_row("SELECT value FROM settings WHERE key = ?1", [key], |r| r.get(0)).optional()?;
713. Ok(match raw {
714. Some(s) => Some(serde_json::from_str(&s)?),
715. None => None,
716. })
717. }
719. pub fn set_setting<T: Serialize>(&self, key: &str, value: &T) -> Result<()> {
720. self.conn.execute(
721. "INSERT INTO settings (key, value) VALUES (?1, ?2) ON CONFLICT (key) DO UPDATE SET value = excluded.value",
722. params![key, serde_json::to_string(value)?],
723. )?;
724. Ok(())
725. }
727. // ---- models and prices ----
729. pub fn replace_models(&self, provider_id: &str, models: &[ModelInfo], source: &str, fetched_at: i64) -> Result<()> {
730. let tx = self.conn.unchecked_transaction()?;
731. tx.execute("DELETE FROM model_pricing WHERE provider_id = ?1", [provider_id])?;
732. for m in models {
733. tx.execute(
734. "INSERT INTO model_pricing (provider_id, model_id, input_per_m, output_per_m, cache_read_per_m,
735. context_length, max_tokens, tags, source, fetched_at)
736. VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
737. params![
738. provider_id,
739. m.id,
740. m.prices.input_per_m,
741. m.prices.output_per_m,
742. m.prices.cache_read_per_m,
743. m.context_length,
744. m.max_tokens,
745. serde_json::to_string(&m.tags)?,
746. source,
747. fetched_at
748. ],
749. )?;
750. }
751. tx.commit()?;
752. Ok(())
753. }
755. pub fn list_models(&self, provider_id: &str) -> Result<Vec<ModelInfo>> {
756. let mut st = self.conn.prepare(
757. "SELECT model_id, input_per_m, output_per_m, cache_read_per_m, context_length, max_tokens, tags
758. FROM model_pricing WHERE provider_id = ?1 ORDER BY lower(model_id)",
759. )?;
760. let rows = st.query_map([provider_id], |r| {
761. let tags: String = r.get(6)?;
762. Ok(ModelInfo {
763. provider_id: provider_id.to_string(),
764. id: r.get(0)?,
765. prices: Prices { input_per_m: r.get(1)?, output_per_m: r.get(2)?, cache_read_per_m: r.get(3)? },
766. context_length: r.get(4)?,
767. max_tokens: r.get(5)?,
768. tags: serde_json::from_str(&tags).unwrap_or_default(),
769. })
770. })?;
771. Ok(rows.collect::<rusqlite::Result<_>>()?)
772. }
774. pub fn model(&self, provider_id: &str, model_id: &str) -> Result<Option<ModelInfo>> {
775. Ok(self.list_models(provider_id)?.into_iter().find(|m| m.id == model_id))
776. }
778. /// `(source, fetched_at)` of a provider's current list, if any.
779. pub fn models_source(&self, provider_id: &str) -> Result<Option<(String, i64)>> {
780. Ok(self
781. .conn
782. .query_row(
783. "SELECT source, MAX(fetched_at) FROM model_pricing WHERE provider_id = ?1 GROUP BY source ORDER BY 2 DESC LIMIT 1",
784. [provider_id],
785. |r| Ok((r.get(0)?, r.get(1)?)),
786. )
787. .optional()?)
788. }
790. // ---- spend ----
792. pub fn spend_since(&self, since_ms: i64) -> Result<(i64, bool)> {
793. Ok(self.conn.query_row(
794. "SELECT COALESCE(SUM(cost_nanousd), 0), COALESCE(MAX(usage_source = 'estimated'), 0)
795. FROM messages WHERE role = 'assistant' AND created_at >= ?1",
796. [since_ms],
797. |r| Ok((r.get(0)?, r.get(1)?)),
798. )?)
799. }
801. pub fn usage_by_model(&self, since_ms: i64) -> Result<Vec<UsageRow>> {
802. self.usage_grouped("model_id", since_ms)
803. }
805. /// Per local day: `utc_offset_minutes` is the viewer's offset from UTC.
806. pub fn usage_by_day(&self, since_ms: i64, utc_offset_minutes: i32) -> Result<Vec<UsageRow>> {
807. let expr = format!("strftime('%Y-%m-%d', (created_at / 1000) + {} * 60, 'unixepoch')", utc_offset_minutes);
808. self.usage_grouped(&expr, since_ms)
809. }
811. fn usage_grouped(&self, key_expr: &str, since_ms: i64) -> Result<Vec<UsageRow>> {
812. let mut st = self.conn.prepare(&format!(
813. "SELECT {key_expr} AS k, SUM(cost_nanousd), SUM(prompt_tokens), SUM(completion_tokens), COUNT(*),
814. MAX(usage_source = 'estimated')
815. FROM messages WHERE role = 'assistant' AND created_at >= ?1 AND k IS NOT NULL
816. GROUP BY k ORDER BY 2 DESC, k"
817. ))?;
818. let rows = st.query_map([since_ms], |r| {
819. Ok(UsageRow {
820. key: r.get(0)?,
821. cost_nanousd: r.get(1)?,
822. prompt_tokens: r.get(2)?,
823. completion_tokens: r.get(3)?,
824. messages: r.get(4)?,
825. estimated: r.get(5)?,
826. })
827. })?;
828. Ok(rows.collect::<rusqlite::Result<_>>()?)
829. }
831. // ---- search ----
833. pub fn search(&self, input: &str, limit: i64) -> Result<Vec<SearchHit>> {
834. let Some(q) = fts_query(input) else { return Ok(vec![]) };
835. let mut st = self.conn.prepare(
836. "SELECT conversation_id, title, message_id, snippet, created_at FROM (
837. SELECT c.id AS conversation_id, c.title AS title, m.id AS message_id,
838. snippet(messages_fts, 0, char(1), char(2), '...', 12) AS snippet,
839. m.created_at AS created_at, bm25(messages_fts) AS rank
840. FROM messages_fts JOIN messages m ON m.rowid = messages_fts.rowid
841. JOIN conversations c ON c.id = m.conversation_id
842. WHERE messages_fts MATCH ?1 AND c.archived = 0
843. UNION ALL
844. SELECT c.id, c.title, NULL, snippet(conversations_fts, 0, char(1), char(2), '...', 12),
845. c.updated_at, bm25(conversations_fts) - 5.0
846. FROM conversations_fts JOIN conversations c ON c.rowid = conversations_fts.rowid
847. WHERE conversations_fts MATCH ?1 AND c.archived = 0
848. ) ORDER BY rank LIMIT ?2",
849. )?;
850. let rows = st.query_map(params![q, limit], |r| {
851. Ok(SearchHit {
852. conversation_id: r.get(0)?,
853. title: r.get(1)?,
854. message_id: r.get(2)?,
855. snippet: r.get(3)?,
856. created_at: r.get(4)?,
857. })
858. })?;
859. Ok(rows.collect::<rusqlite::Result<_>>()?)
860. }
862. #[cfg(test)]
863. pub(crate) fn conn(&self) -> &Connection {
864. &self.conn
865. }
866. }
868. /// The first line of the first message, trimmed to a sensible length.
869. pub fn title_from(content: &str) -> String {
870. let line = content.lines().map(str::trim).find(|l| !l.is_empty()).unwrap_or("New chat");
871. let mut t: String = line.chars().take(60).collect();
872. if line.chars().count() > 60 {
873. if let Some(cut) = t.rfind(' ').filter(|&i| i > 30) {
874. t.truncate(cut);
875. }
876. t.push_str("...");
877. }
878. t
879. }
881. #[cfg(test)]
882. mod tests;