turn.rs

crates/cortex-core/tests/turn.rs

1. //! Whole turns, database included, against the stand-in server.
3. use std::sync::Mutex;
5. use cortex_core::api::StreamEvent;
6. use cortex_core::chat::{build_request, run_turn};
7. use cortex_core::pricing::parse_priced_models;
8. use cortex_core::provider::Client;
9. use cortex_core::storage::{Db, NewAttachment, DEEPINFRA_ID};
10. use cortex_core::types::UsageSource;
11. use tokio_util::sync::CancellationToken;
13. struct Setup {
14. db: Mutex<Db>,
15. client: Client,
16. _s: standin::Standin,
17. }
19. async fn setup() -> Setup {
20. let s = standin::spawn().await;
21. let db = Db::open_in_memory().unwrap();
22. db.replace_models(DEEPINFRA_ID, &parse_priced_models(DEEPINFRA_ID, &standin::models()), "live", 0).unwrap();
23. let client = Client::new(&s.base_url(), Some("good-key".into()));
24. Setup { db: Mutex::new(db), client, _s: s }
25. }
27. async fn turn(
28. st: &Setup,
29. model: &str,
30. text: &str,
31. cancel: &CancellationToken,
32. ) -> (cortex_core::storage::Message, Vec<StreamEvent>) {
33. let (req, assistant, prices) = {
34. let db = st.db.lock().unwrap();
35. let conv = db.create_conversation(DEEPINFRA_ID, model, "Answer briefly.", None, "main").unwrap();
36. db.insert_user_message(&conv.id, text, &[]).unwrap();
37. let req = build_request(&db, &conv.id).unwrap();
38. let a = db.insert_assistant_placeholder(&conv.id, DEEPINFRA_ID, model).unwrap();
39. let prices = db.model(DEEPINFRA_ID, model).unwrap().map(|m| m.prices);
40. (req, a, prices)
41. };
42. let mut events = Vec::new();
43. let m = run_turn(&st.db, &st.client, &req, &assistant, prices, cancel, |e| events.push(e)).await.unwrap();
44. (m, events)
45. }
47. #[tokio::test]
48. async fn a_priced_reply_is_saved_with_its_cost() {
49. let st = setup().await;
50. let (m, events) = turn(&st, "standin/echo", "tidy my backup script", &CancellationToken::new()).await;
51. assert_eq!(m.status, "complete");
52. assert_eq!(m.content, "You said: tidy my backup script");
53. assert_eq!(m.usage_source, UsageSource::Reported);
54. // system (2 words) + user (4) + 4 per message x2 = 14 in; 6 words out.
55. assert_eq!((m.usage.prompt_tokens, m.usage.completion_tokens), (14, 6));
56. // 14 x $0.50/M + 6 x $1.50/M = 7,000 + 9,000 nano
57. assert_eq!(m.cost_nanousd, 16_000);
58. assert_eq!(m.prices.unwrap().input_per_m, 0.5);
59. let text: String = events
60. .iter()
61. .filter_map(|e| if let StreamEvent::Delta { text } = e { Some(text.as_str()) } else { None })
62. .collect();
63. assert_eq!(text, m.content);
64. assert!(matches!(events.last(), Some(StreamEvent::Done { .. })));
65. assert!(events.iter().any(|e| matches!(e, StreamEvent::Usage { .. })));
66. }
68. #[tokio::test]
69. async fn no_usage_means_an_estimate_marked_as_such() {
70. let st = setup().await;
71. let (m, _) = turn(&st, "standin/no-usage", "hello", &CancellationToken::new()).await;
72. assert_eq!(m.usage_source, UsageSource::Estimated);
73. assert!(m.usage.prompt_tokens > 0 && m.usage.completion_tokens > 0);
74. assert!(m.cost_nanousd > 0);
75. }
77. #[tokio::test]
78. async fn cached_prompt_tokens_bill_at_the_cache_price() {
79. let st = setup().await;
80. let (m, _) =
81. turn(&st, "standin/cached", "one two three four five six seven eight nine ten", &CancellationToken::new())
82. .await;
83. // 2 + 10 words + 4 per message = 20 in, all of it cached (the stand-in caches up to 40).
84. let u = m.usage;
85. assert_eq!((u.prompt_tokens, u.cached_tokens, u.completion_tokens), (20, 20, 12));
86. let expect = (u.prompt_tokens - u.cached_tokens) * 500 + u.cached_tokens * 100 + u.completion_tokens * 1500;
87. assert_eq!(m.cost_nanousd, expect);
88. }
90. #[tokio::test]
91. async fn cancelling_saves_the_partial_reply() {
92. let st = setup().await;
93. let cancel = CancellationToken::new();
94. let c2 = cancel.clone();
95. tokio::spawn(async move {
96. tokio::time::sleep(std::time::Duration::from_millis(400)).await;
97. c2.cancel();
98. });
99. let (m, events) = turn(&st, "standin/slow", "go", &cancel).await;
100. assert_eq!(m.status, "cancelled");
101. assert!(m.content.starts_with("word1"), "{}", m.content);
102. assert!(!m.content.contains("word200"));
103. assert_eq!(m.usage_source, UsageSource::Estimated);
104. let saved = st.db.lock().unwrap().get_message(&m.id).unwrap();
105. assert_eq!(saved.status, "cancelled");
106. assert!(matches!(events.last(), Some(StreamEvent::Done { .. })));
107. }
109. #[tokio::test]
110. async fn failures_are_saved_and_reported() {
111. let st = setup().await;
112. let (m, events) = turn(&st, "standin/402", "x", &CancellationToken::new()).await;
113. assert_eq!((m.status.as_str(), m.error_kind.as_deref()), ("error", Some("insufficient_balance")));
114. assert_eq!((m.cost_nanousd, m.usage_source), (0, UsageSource::None));
115. assert!(matches!(events.last(), Some(StreamEvent::Error { kind, .. }) if kind == "insufficient_balance"));
117. let (m, _) = turn(&st, "standin/midstream-error", "a b c d", &CancellationToken::new()).await;
118. assert_eq!(m.status, "error");
119. assert_eq!(m.content, "You said: a");
120. assert_eq!(m.usage_source, UsageSource::Estimated);
121. }
123. #[tokio::test]
124. async fn history_attachments_and_failed_replies() {
125. let st = setup().await;
126. let dir = tempfile::tempdir().unwrap();
127. let img_path = dir.path().join("x.png");
128. std::fs::write(&img_path, b"PNGDATA").unwrap();
129. let req = {
130. let db = st.db.lock().unwrap();
131. let conv = db.create_conversation(DEEPINFRA_ID, "standin/vision", "", None, "main").unwrap();
132. let img = db
133. .insert_attachment(&NewAttachment {
134. kind: "image",
135. filename: "chart.png",
136. mime: "image/png",
137. byte_size: 7,
138. sha256: "h",
139. width: Some(1),
140. height: Some(1),
141. path: Some(img_path.to_str().unwrap()),
142. text_content: None,
143. })
144. .unwrap();
145. let txt = db
146. .insert_attachment(&NewAttachment {
147. kind: "text",
148. filename: "notes.md",
149. mime: "text/markdown",
150. byte_size: 5,
151. sha256: "t",
152. width: None,
153. height: None,
154. path: None,
155. text_content: Some("Q3 was flat."),
156. })
157. .unwrap();
158. db.insert_user_message(&conv.id, "compare these", &[img.id, txt.id]).unwrap();
159. let bad = db.insert_assistant_placeholder(&conv.id, DEEPINFRA_ID, "standin/vision").unwrap();
160. db.finalize_assistant(
161. &bad.id,
162. &cortex_core::storage::Finalize { content: "half".into(), status: "error", ..Default::default() },
163. )
164. .unwrap();
165. db.insert_user_message(&conv.id, "try again", &[]).unwrap();
166. build_request(&db, &conv.id).unwrap()
167. };
168. let v = serde_json::to_value(&req).unwrap();
169. let msgs = v["messages"].as_array().unwrap();
170. assert_eq!(msgs.len(), 2, "no system prompt, and the failed reply is left out: {v}");
171. let parts = msgs[0]["content"].as_array().unwrap();
172. assert_eq!(parts[0]["text"], "compare these\n\nFile: notes.md\n```\nQ3 was flat.\n```");
173. assert_eq!(parts[1]["image_url"]["url"], "data:image/png;base64,UE5HREFUQQ==");
174. assert_eq!(msgs[1]["content"], "try again");
175. // And the vision stand-in counts the one image.
176. let a = {
177. let db = st.db.lock().unwrap();
178. let conv_id = db.list_conversations(false, 1).unwrap()[0].id.clone();
179. db.insert_assistant_placeholder(&conv_id, DEEPINFRA_ID, "standin/vision").unwrap()
180. };
181. let m = run_turn(&st.db, &st.client, &req, &a, None, &CancellationToken::new(), |_| {}).await.unwrap();
182. assert_eq!(m.content, "I see 1 image.");
183. assert_eq!(m.cost_nanousd, 0, "unknown prices cost nothing rather than a guess");
184. }