1use crate::coercion::CoercionFlag;
6use crate::message::Message;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct HealingMetadata {
15 pub flags: Vec<CoercionFlag>,
17 pub confidence: f32,
19 pub original_error: String,
21}
22
23impl HealingMetadata {
24 pub fn new(flags: Vec<CoercionFlag>, confidence: f32, original_error: String) -> Self {
26 Self {
27 flags,
28 confidence: confidence.clamp(0.0, 1.0),
29 original_error,
30 }
31 }
32
33 pub fn has_major_coercions(&self) -> bool {
35 self.flags.iter().any(|f| f.is_major())
36 }
37
38 pub fn is_confident(&self, threshold: f32) -> bool {
40 self.confidence >= threshold
41 }
42}
43
44#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
46pub struct CompletionResponse {
47 pub id: String,
49 pub model: String,
51 pub choices: Vec<CompletionChoice>,
53 pub usage: Usage,
55 #[serde(skip_serializing_if = "Option::is_none")]
57 pub created: Option<i64>,
58 #[serde(skip_serializing_if = "Option::is_none")]
60 pub provider: Option<String>,
61 #[serde(skip_serializing_if = "Option::is_none")]
63 pub healing_metadata: Option<HealingMetadata>,
64}
65
66impl CompletionResponse {
67 pub fn content(&self) -> Option<&str> {
96 self.choices
97 .first()
98 .map(|choice| choice.message.content.as_str())
99 }
100
101 pub fn first_choice(&self) -> Option<&CompletionChoice> {
103 self.choices.first()
104 }
105
106 pub fn was_healed(&self) -> bool {
138 self.healing_metadata.is_some()
139 }
140
141 pub fn confidence(&self) -> f32 {
173 self.healing_metadata
174 .as_ref()
175 .map(|m| m.confidence)
176 .unwrap_or(1.0)
177 }
178}
179
180#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
182pub struct CompletionChoice {
183 pub index: u32,
185 pub message: Message,
187 pub finish_reason: FinishReason,
189 #[serde(skip_serializing_if = "Option::is_none")]
191 pub logprobs: Option<serde_json::Value>,
192}
193
194#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
196#[serde(rename_all = "snake_case")]
197pub enum FinishReason {
198 Stop,
200 Length,
202 ContentFilter,
204 ToolCalls,
206}
207
208impl FinishReason {
209 pub fn as_str(self) -> &'static str {
211 match self {
212 Self::Stop => "stop",
213 Self::Length => "length",
214 Self::ContentFilter => "content_filter",
215 Self::ToolCalls => "tool_calls",
216 }
217 }
218}
219
220#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
222pub struct Usage {
223 pub prompt_tokens: u32,
225 pub completion_tokens: u32,
227 pub total_tokens: u32,
229}
230
231impl Usage {
232 pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
234 Self {
235 prompt_tokens,
236 completion_tokens,
237 total_tokens: prompt_tokens + completion_tokens,
238 }
239 }
240}
241
242#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
244pub struct CompletionChunk {
245 pub id: String,
247 pub model: String,
249 pub choices: Vec<ChoiceDelta>,
251 #[serde(skip_serializing_if = "Option::is_none")]
253 pub created: Option<i64>,
254 #[serde(skip_serializing_if = "Option::is_none")]
256 pub usage: Option<Usage>,
257}
258
259#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
261pub struct ChoiceDelta {
262 pub index: u32,
264 pub delta: MessageDelta,
266 #[serde(skip_serializing_if = "Option::is_none")]
268 pub finish_reason: Option<FinishReason>,
269}
270
271#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
273pub struct MessageDelta {
274 #[serde(skip_serializing_if = "Option::is_none")]
276 pub role: Option<crate::message::Role>,
277 #[serde(skip_serializing_if = "Option::is_none")]
279 pub content: Option<String>,
280 #[serde(skip_serializing_if = "Option::is_none")]
282 pub reasoning_content: Option<String>,
283 #[serde(skip_serializing_if = "Option::is_none")]
285 pub tool_calls: Option<Vec<ToolCallDelta>>,
286}
287
288#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
290pub struct ToolCallDelta {
291 pub index: u32,
293 #[serde(skip_serializing_if = "Option::is_none")]
295 pub id: Option<String>,
296 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
298 pub tool_type: Option<crate::tool::ToolType>,
299 #[serde(skip_serializing_if = "Option::is_none")]
301 pub function: Option<ToolCallFunctionDelta>,
302}
303
304#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
306pub struct ToolCallFunctionDelta {
307 #[serde(skip_serializing_if = "Option::is_none")]
309 pub name: Option<String>,
310 #[serde(skip_serializing_if = "Option::is_none")]
312 pub arguments: Option<String>,
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318
319 #[test]
320 fn test_completion_response_content() {
321 let response = CompletionResponse {
322 id: "resp_123".to_string(),
323 model: "gpt-4".to_string(),
324 choices: vec![CompletionChoice {
325 index: 0,
326 message: Message::assistant("Hello!"),
327 finish_reason: FinishReason::Stop,
328 logprobs: None,
329 }],
330 usage: Usage::new(10, 5),
331 created: Some(1234567890),
332 provider: Some("openai".to_string()),
333 healing_metadata: None,
334 };
335
336 assert_eq!(response.content(), Some("Hello!"));
337 assert_eq!(response.first_choice().unwrap().index, 0);
338 assert!(!response.was_healed());
339 assert_eq!(response.confidence(), 1.0);
340 }
341
342 #[test]
343 fn test_completion_response_empty_choices() {
344 let response = CompletionResponse {
345 id: "resp_123".to_string(),
346 model: "gpt-4".to_string(),
347 choices: vec![],
348 usage: Usage::new(10, 0),
349 created: None,
350 provider: None,
351 healing_metadata: None,
352 };
353
354 assert_eq!(response.content(), None);
355 assert_eq!(response.first_choice(), None);
356 }
357
358 #[test]
359 fn test_usage_calculation() {
360 let usage = Usage::new(100, 50);
361 assert_eq!(usage.prompt_tokens, 100);
362 assert_eq!(usage.completion_tokens, 50);
363 assert_eq!(usage.total_tokens, 150);
364 }
365
366 #[test]
367 fn test_finish_reason_serialization() {
368 let json = serde_json::to_string(&FinishReason::Stop).unwrap();
369 assert_eq!(json, "\"stop\"");
370
371 let json = serde_json::to_string(&FinishReason::Length).unwrap();
372 assert_eq!(json, "\"length\"");
373
374 let json = serde_json::to_string(&FinishReason::ContentFilter).unwrap();
375 assert_eq!(json, "\"content_filter\"");
376
377 let json = serde_json::to_string(&FinishReason::ToolCalls).unwrap();
378 assert_eq!(json, "\"tool_calls\"");
379 }
380
381 #[test]
382 fn test_response_serialization() {
383 let response = CompletionResponse {
384 id: "resp_123".to_string(),
385 model: "gpt-4".to_string(),
386 choices: vec![CompletionChoice {
387 index: 0,
388 message: Message::assistant("Hello!"),
389 finish_reason: FinishReason::Stop,
390 logprobs: None,
391 }],
392 usage: Usage::new(10, 5),
393 created: None,
394 provider: None,
395 healing_metadata: None,
396 };
397
398 let json = serde_json::to_string(&response).unwrap();
399 let parsed: CompletionResponse = serde_json::from_str(&json).unwrap();
400 assert_eq!(response, parsed);
401 }
402
403 #[test]
404 fn test_streaming_chunk() {
405 let chunk = CompletionChunk {
406 id: "resp_123".to_string(),
407 model: "gpt-4".to_string(),
408 choices: vec![ChoiceDelta {
409 index: 0,
410 delta: MessageDelta {
411 role: Some(crate::message::Role::Assistant),
412 content: Some("Hello".to_string()),
413 reasoning_content: None,
414 tool_calls: None,
415 },
416 finish_reason: None,
417 }],
418 created: Some(1234567890),
419 usage: None,
420 };
421
422 let json = serde_json::to_string(&chunk).unwrap();
423 let parsed: CompletionChunk = serde_json::from_str(&json).unwrap();
424 assert_eq!(chunk, parsed);
425 }
426
427 #[test]
428 fn test_message_delta() {
429 let delta = MessageDelta {
430 role: Some(crate::message::Role::Assistant),
431 content: Some("Hi".to_string()),
432 reasoning_content: None,
433 tool_calls: None,
434 };
435
436 let json = serde_json::to_value(&delta).unwrap();
437 assert_eq!(json.get("role").and_then(|v| v.as_str()), Some("assistant"));
438 assert_eq!(json.get("content").and_then(|v| v.as_str()), Some("Hi"));
439 }
440
441 #[test]
442 fn test_optional_fields_not_serialized() {
443 let response = CompletionResponse {
444 id: "resp_123".to_string(),
445 model: "gpt-4".to_string(),
446 choices: vec![],
447 usage: Usage::new(10, 5),
448 created: None,
449 provider: None,
450 healing_metadata: None,
451 };
452
453 let json = serde_json::to_value(&response).unwrap();
454 assert!(json.get("created").is_none());
455 assert!(json.get("provider").is_none());
456 assert!(json.get("healing_metadata").is_none());
457 }
458
459 #[test]
460 fn test_healing_metadata() {
461 use crate::coercion::CoercionFlag;
462
463 let metadata = HealingMetadata::new(
464 vec![CoercionFlag::StrippedMarkdown],
465 0.9,
466 "Parse error".to_string(),
467 );
468
469 assert_eq!(metadata.confidence, 0.9);
470 assert!(!metadata.has_major_coercions());
471 assert!(metadata.is_confident(0.8));
472 assert!(!metadata.is_confident(0.95));
473
474 let major_metadata = HealingMetadata::new(
475 vec![CoercionFlag::TruncatedJson],
476 0.7,
477 "Parse error".to_string(),
478 );
479
480 assert!(major_metadata.has_major_coercions());
481 }
482
483 #[test]
484 fn test_healing_metadata_confidence_clamped() {
485 let metadata = HealingMetadata::new(vec![], 1.5, "error".to_string());
486 assert_eq!(metadata.confidence, 1.0);
487
488 let metadata = HealingMetadata::new(vec![], -0.5, "error".to_string());
489 assert_eq!(metadata.confidence, 0.0);
490 }
491
492 #[test]
493 fn test_response_with_healing_metadata() {
494 use crate::coercion::CoercionFlag;
495
496 let metadata = HealingMetadata::new(
497 vec![
498 CoercionFlag::StrippedMarkdown,
499 CoercionFlag::FixedTrailingComma,
500 ],
501 0.85,
502 "JSON parse error".to_string(),
503 );
504
505 let response = CompletionResponse {
506 id: "resp_123".to_string(),
507 model: "gpt-4".to_string(),
508 choices: vec![],
509 usage: Usage::new(10, 5),
510 created: None,
511 provider: None,
512 healing_metadata: Some(metadata),
513 };
514
515 assert!(response.was_healed());
516 assert_eq!(response.confidence(), 0.85);
517
518 let json = serde_json::to_string(&response).unwrap();
519 let parsed: CompletionResponse = serde_json::from_str(&json).unwrap();
520 assert_eq!(response, parsed);
521 assert!(parsed.was_healed());
522 assert_eq!(parsed.confidence(), 0.85);
523 }
524}