chat.rs

crates/cortex-core/src/chat.rs

1. //! One turn of a conversation: build the request from history, stream the
2. //! reply, and save it with its tokens and cost. The Tauri command is a thin
3. //! wrapper over `run_turn`, so the whole path is tested without a window.
5. use std::sync::{Mutex, PoisonError};
6. use std::time::{Duration, Instant};
8. use tokio_util::sync::CancellationToken;
10. use crate::api::StreamEvent;
11. use crate::attachments::{data_url, text_block};
12. use crate::cost::cost_nanousd;
13. use crate::estimate::{estimate_prompt, estimate_tokens};
14. use crate::provider::openai::{OutContent, Part};
15. use crate::provider::{ChatRequest, Client, OutMessage, StreamItem, StreamOptions};
16. use crate::storage::{Db, DbError, Finalize, Message};
17. use crate::types::{Prices, Usage, UsageSource};
19. /// Deltas are batched to about this often, so the interface is not flooded.
20. pub const COALESCE: Duration = Duration::from_millis(30);
22. /// The request for a conversation as it stands: system prompt, then every
23. /// finished message, with text files inlined and images as data URLs.
24. /// Replies that failed or were stopped are left out of the model's view.
25. pub fn build_request(db: &Db, conversation_id: &str) -> Result<ChatRequest, DbError> {
26. let conv = db.get_conversation(conversation_id)?;
27. let preset = match &conv.preset_id {
28. Some(id) => db.list_presets()?.into_iter().find(|p| &p.id == id),
29. None => None,
30. };
31. let mut messages = Vec::new();
32. if !conv.system_prompt.trim().is_empty() {
33. messages.push(OutMessage::text("system", conv.system_prompt.trim()));
34. }
35. for m in db.list_messages(conversation_id)? {
36. match m.role.as_str() {
37. "user" => {
38. let mut text = m.content.clone();
39. let mut images = Vec::new();
40. for a in db.attachment_data_for(&m.id)? {
41. match a.meta.kind.as_str() {
42. "text" => {
43. if !text.is_empty() {
44. text.push_str("\n\n");
45. }
46. text.push_str(&text_block(&a.meta.filename, a.text_content.as_deref().unwrap_or("")));
47. }
48. "image" => {
49. if let Some(bytes) = a.path.as_deref().and_then(|p| std::fs::read(p).ok()) {
50. images.push(data_url(&a.meta.mime, &bytes));
51. }
52. }
53. _ => {}
54. }
55. }
56. messages.push(OutMessage::with_images("user", text, images));
57. }
58. "assistant" if m.status == "complete" && !m.content.is_empty() => {
59. messages.push(OutMessage::text("assistant", m.content.clone()));
60. }
61. _ => {}
62. }
63. }
64. Ok(ChatRequest {
65. model: conv.model_id,
66. messages,
67. stream: true,
68. stream_options: Some(StreamOptions { include_usage: true }),
69. temperature: preset.as_ref().and_then(|p| p.temperature),
70. max_tokens: preset.as_ref().and_then(|p| p.max_tokens),
71. })
72. }
74. /// Rough prompt size of a request, for estimates.
75. pub fn request_tokens(req: &ChatRequest) -> i64 {
76. let msgs: Vec<(String, usize)> = req
77. .messages
78. .iter()
79. .map(|m| match &m.content {
80. OutContent::Text(s) => (s.clone(), 0),
81. OutContent::Parts(parts) => {
82. let text = parts
83. .iter()
84. .filter_map(|p| match p {
85. Part::Text { text } => Some(text.as_str()),
86. Part::ImageUrl { .. } => None,
87. })
88. .collect::<Vec<_>>()
89. .join("\n");
90. let images = parts.iter().filter(|p| matches!(p, Part::ImageUrl { .. })).count();
91. (text, images)
92. }
93. })
94. .collect();
95. estimate_prompt(msgs.iter().map(|(t, n)| (t.as_str(), *n)))
96. }
98. /// Batches deltas: flushes when `COALESCE` has passed since the last flush.
99. struct Coalescer {
100. content: String,
101. reasoning: String,
102. last: Instant,
103. }
105. impl Coalescer {
106. fn new() -> Self {
107. Self { content: String::new(), reasoning: String::new(), last: Instant::now() }
108. }
109. fn flush(&mut self, emit: &mut impl FnMut(StreamEvent)) {
110. if !self.reasoning.is_empty() {
111. emit(StreamEvent::Reasoning { text: std::mem::take(&mut self.reasoning) });
112. }
113. if !self.content.is_empty() {
114. emit(StreamEvent::Delta { text: std::mem::take(&mut self.content) });
115. }
116. self.last = Instant::now();
117. }
118. fn push(&mut self, item: StreamItem, emit: &mut impl FnMut(StreamEvent)) {
119. match item {
120. StreamItem::Content(t) => self.content.push_str(&t),
121. StreamItem::Reasoning(t) => self.reasoning.push_str(&t),
122. StreamItem::Usage(u) => {
123. self.flush(emit);
124. emit(StreamEvent::Usage { usage: u });
125. return;
126. }
127. }
128. if self.last.elapsed() >= COALESCE {
129. self.flush(emit);
130. }
131. }
132. }
134. /// Stream the reply into the placeholder `assistant` row and save the outcome.
135. /// Always returns the saved row: complete, cancelled (keeping what arrived), or error.
136. pub async fn run_turn(
137. db: &Mutex<Db>,
138. client: &Client,
139. req: &ChatRequest,
140. assistant: &Message,
141. prices: Option<Prices>,
142. cancel: &CancellationToken,
143. mut emit: impl FnMut(StreamEvent),
144. ) -> Result<Message, DbError> {
145. let mut co = Coalescer::new();
146. let out = client.stream_chat(req, cancel, |item| co.push(item, &mut emit)).await;
147. co.flush(&mut emit);
149. let got_anything = !out.content.is_empty() || !out.reasoning.is_empty();
150. let (usage, source) = match (out.usage, got_anything) {
151. (Some(u), _) => (u, UsageSource::Reported),
152. // Nothing reported: estimate from what was sent and what came back.
153. (None, true) => (
154. Usage {
155. prompt_tokens: request_tokens(req),
156. completion_tokens: estimate_tokens(&out.content) + estimate_tokens(&out.reasoning),
157. cached_tokens: 0,
158. },
159. UsageSource::Estimated,
160. ),
161. // Refused before a single token: nothing billed.
162. (None, false) => (Usage::default(), UsageSource::None),
163. };
164. let cost = prices.map(|p| cost_nanousd(&usage, &p)).unwrap_or(0);
165. let status = if out.cancelled {
166. "cancelled"
167. } else if out.error.is_some() {
168. "error"
169. } else {
170. "complete"
171. };
172. let f = Finalize {
173. content: out.content.clone(),
174. reasoning: Some(out.reasoning.clone()),
175. status,
176. error_kind: out.error.as_ref().map(|e| e.kind().to_string()),
177. error_message: out.error.as_ref().map(|e| e.to_string()),
178. usage,
179. usage_source: Some(source),
180. prices,
181. cost_nanousd: cost,
182. provider_cost_nanousd: out.provider_cost_usd.map(|c| (c * crate::cost::NANO_PER_USD).round() as i64),
183. finish_reason: out.finish_reason.clone(),
184. ttft_ms: out.ttft_ms.map(|v| v as i64),
185. duration_ms: Some(out.duration_ms as i64),
186. };
187. let saved = db.lock().unwrap_or_else(PoisonError::into_inner).finalize_assistant(&assistant.id, &f)?;
188. match &out.error {
189. Some(e) => emit(StreamEvent::Error { kind: e.kind().into(), error: e.to_string(), message: saved.clone() }),
190. None => emit(StreamEvent::Done { message: saved.clone() }),
191. }
192. Ok(saved)
193. }