openai.rs

crates/cortex-core/src/provider/openai.rs

1. //! The OpenAI-compatible chat wire format, as DeepInfra, Ollama, LM Studio and
2. //! llama.cpp speak it. Parsing is deliberately tolerant: unknown fields are
3. //! ignored and every field is optional, because each server differs a little.
5. use serde::{Deserialize, Serialize};
6. use serde_json::Value;
8. use super::error::{error_message, ProviderError};
9. use crate::types::Usage;
11. #[derive(Debug, Clone, Serialize, PartialEq)]
12. pub struct ChatRequest {
13. pub model: String,
14. pub messages: Vec<OutMessage>,
15. pub stream: bool,
16. #[serde(skip_serializing_if = "Option::is_none")]
17. pub stream_options: Option<StreamOptions>,
18. #[serde(skip_serializing_if = "Option::is_none")]
19. pub temperature: Option<f32>,
20. #[serde(skip_serializing_if = "Option::is_none")]
21. pub max_tokens: Option<u32>,
22. }
24. #[derive(Debug, Clone, Serialize, PartialEq)]
25. pub struct StreamOptions {
26. pub include_usage: bool,
27. }
29. #[derive(Debug, Clone, Serialize, PartialEq)]
30. pub struct OutMessage {
31. pub role: String,
32. pub content: OutContent,
33. }
35. #[derive(Debug, Clone, Serialize, PartialEq)]
36. #[serde(untagged)]
37. pub enum OutContent {
38. Text(String),
39. Parts(Vec<Part>),
40. }
42. #[derive(Debug, Clone, Serialize, PartialEq)]
43. #[serde(tag = "type", rename_all = "snake_case")]
44. pub enum Part {
45. Text { text: String },
46. ImageUrl { image_url: ImageUrl },
47. }
49. #[derive(Debug, Clone, Serialize, PartialEq)]
50. pub struct ImageUrl {
51. pub url: String,
52. }
54. impl OutMessage {
55. pub fn text(role: &str, text: impl Into<String>) -> Self {
56. Self { role: role.into(), content: OutContent::Text(text.into()) }
57. }
59. /// Text plus images (data: URLs). Plain text when there are no images, since
60. /// some local servers reject the parts form entirely.
61. pub fn with_images(role: &str, text: impl Into<String>, image_urls: Vec<String>) -> Self {
62. let text = text.into();
63. if image_urls.is_empty() {
64. return Self::text(role, text);
65. }
66. let mut parts = vec![Part::Text { text }];
67. parts.extend(image_urls.into_iter().map(|url| Part::ImageUrl { image_url: ImageUrl { url } }));
68. Self { role: role.into(), content: OutContent::Parts(parts) }
69. }
70. }
72. /// What one `data:` payload meant.
73. #[derive(Debug, Clone, PartialEq, Default)]
74. pub struct Chunk {
75. pub content: Option<String>,
76. pub reasoning: Option<String>,
77. pub finish_reason: Option<String>,
78. pub usage: Option<Usage>,
79. /// DeepInfra may report its own cost with the usage; kept for comparison only.
80. pub provider_cost_usd: Option<f64>,
81. pub done: bool,
82. }
84. #[derive(Deserialize)]
85. struct RawChunk {
86. #[serde(default)]
87. choices: Vec<RawChoice>,
88. #[serde(default)]
89. usage: Option<RawUsage>,
90. }
92. #[derive(Deserialize)]
93. struct RawChoice {
94. #[serde(default)]
95. delta: Option<RawDelta>,
96. #[serde(default)]
97. finish_reason: Option<String>,
98. }
100. #[derive(Deserialize, Default)]
101. struct RawDelta {
102. #[serde(default)]
103. content: Option<String>,
104. #[serde(default)]
105. reasoning_content: Option<String>,
106. /// Ollama and some proxies use `reasoning`.
107. #[serde(default)]
108. reasoning: Option<String>,
109. }
111. #[derive(Deserialize)]
112. struct RawUsage {
113. #[serde(default)]
114. prompt_tokens: Option<i64>,
115. #[serde(default)]
116. completion_tokens: Option<i64>,
117. #[serde(default)]
118. prompt_tokens_details: Option<Value>,
119. #[serde(default)]
120. estimated_cost: Option<f64>,
121. }
123. pub fn parse_chunk(data: &str) -> Result<Chunk, ProviderError> {
124. let data = data.trim();
125. if data == "[DONE]" {
126. return Ok(Chunk { done: true, ..Default::default() });
127. }
128. let v: Value = serde_json::from_str(data).map_err(|e| ProviderError::Stream(format!("unreadable chunk: {e}")))?;
129. if v.get("error").is_some() {
130. return Err(ProviderError::Stream(error_message(data).unwrap_or_else(|| "unknown error".into())));
131. }
132. let raw: RawChunk =
133. serde_json::from_value(v).map_err(|e| ProviderError::Stream(format!("unexpected chunk: {e}")))?;
134. let mut out = Chunk::default();
135. if let Some(choice) = raw.choices.into_iter().next() {
136. let d = choice.delta.unwrap_or_default();
137. out.content = d.content.filter(|s| !s.is_empty());
138. out.reasoning = d.reasoning_content.or(d.reasoning).filter(|s| !s.is_empty());
139. out.finish_reason = choice.finish_reason;
140. }
141. if let Some(u) = raw.usage {
142. if u.prompt_tokens.is_some() || u.completion_tokens.is_some() {
143. let cached = u
144. .prompt_tokens_details
145. .as_ref()
146. .and_then(|d| d.get("cached_tokens"))
147. .and_then(Value::as_i64)
148. .unwrap_or(0);
149. out.usage = Some(Usage {
150. prompt_tokens: u.prompt_tokens.unwrap_or(0),
151. completion_tokens: u.completion_tokens.unwrap_or(0),
152. cached_tokens: cached,
153. });
154. }
155. out.provider_cost_usd = u.estimated_cost;
156. }
157. Ok(out)
158. }
160. #[cfg(test)]
161. mod tests {
162. use super::*;
164. #[test]
165. fn request_shape() {
166. let req = ChatRequest {
167. model: "m".into(),
168. messages: vec![
169. OutMessage::text("system", "be brief"),
170. OutMessage::with_images("user", "what is this", vec!["data:image/png;base64,AAA".into()]),
171. OutMessage::with_images("user", "plain", vec![]),
172. ],
173. stream: true,
174. stream_options: Some(StreamOptions { include_usage: true }),
175. temperature: None,
176. max_tokens: Some(512),
177. };
178. let v = serde_json::to_value(&req).unwrap();
179. assert_eq!(v["stream_options"]["include_usage"], true);
180. assert_eq!(v["messages"][0]["content"], "be brief");
181. assert_eq!(v["messages"][1]["content"][0], serde_json::json!({"type": "text", "text": "what is this"}));
182. assert_eq!(v["messages"][1]["content"][1]["image_url"]["url"], "data:image/png;base64,AAA");
183. assert_eq!(v["messages"][2]["content"], "plain");
184. assert!(v.get("temperature").is_none());
185. }
187. #[test]
188. fn deepinfra_shape_with_final_usage_chunk() {
189. let c = parse_chunk(r#"{"id":"x","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"},"finish_reason":null}]}"#).unwrap();
190. assert_eq!(c.content.as_deref(), Some("Hel"));
191. let last = parse_chunk(r#"{"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":30,"total_tokens":42,"prompt_tokens_details":{"cached_tokens":8},"estimated_cost":0.0000123}}"#).unwrap();
192. assert_eq!(last.usage, Some(Usage { prompt_tokens: 12, completion_tokens: 30, cached_tokens: 8 }));
193. assert_eq!(last.provider_cost_usd, Some(0.0000123));
194. assert!(parse_chunk("[DONE]").unwrap().done);
195. }
197. #[test]
198. fn usage_on_last_content_chunk_and_finish_reason() {
199. let c = parse_chunk(r#"{"choices":[{"delta":{"content":"!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2}}"#).unwrap();
200. assert_eq!(c.content.as_deref(), Some("!"));
201. assert_eq!(c.finish_reason.as_deref(), Some("stop"));
202. assert_eq!(c.usage.unwrap().completion_tokens, 2);
203. }
205. #[test]
206. fn ollama_and_llamacpp_shapes() {
207. // Ollama: reasoning under `reasoning`, no usage unless asked.
208. let c = parse_chunk(r#"{"id":"chatcmpl-1","model":"qwen3","choices":[{"index":0,"delta":{"role":"assistant","content":"","reasoning":"thinking"},"finish_reason":null}]}"#).unwrap();
209. assert_eq!(c.reasoning.as_deref(), Some("thinking"));
210. assert_eq!(c.content, None);
211. // llama.cpp: usage with nulls elsewhere, timings block ignored.
212. let c = parse_chunk(r#"{"choices":[{"finish_reason":"stop","index":0,"delta":{}}],"usage":{"completion_tokens":9,"prompt_tokens":4,"total_tokens":13},"timings":{"predicted_ms":10}}"#).unwrap();
213. assert_eq!(c.usage.unwrap().prompt_tokens, 4);
214. }
216. #[test]
217. fn deepinfra_reasoning_content() {
218. let c = parse_chunk(r#"{"choices":[{"delta":{"reasoning_content":"step 1"}}]}"#).unwrap();
219. assert_eq!(c.reasoning.as_deref(), Some("step 1"));
220. }
222. #[test]
223. fn errors_in_stream() {
224. let e = parse_chunk(r#"{"error":{"message":"model overloaded"}}"#).unwrap_err();
225. assert_eq!(e, ProviderError::Stream("model overloaded".into()));
226. assert_eq!(parse_chunk("not json").unwrap_err().kind(), "stream");
227. }
228. }