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> {
97 self.choices
98 .first()
99 .map(|choice| choice.message.content.as_str())
100 }
101
102 pub fn first_choice(&self) -> Option<&CompletionChoice> {
104 self.choices.first()
105 }
106
107 pub fn was_healed(&self) -> bool {
139 self.healing_metadata.is_some()
140 }
141
142 pub fn confidence(&self) -> f32 {
174 self.healing_metadata
175 .as_ref()
176 .map(|m| m.confidence)
177 .unwrap_or(1.0)
178 }
179}
180
181#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
183pub struct CompletionChoice {
184 pub index: u32,
186 pub message: Message,
188 pub finish_reason: FinishReason,
190 #[serde(skip_serializing_if = "Option::is_none")]
192 pub logprobs: Option<serde_json::Value>,
193}
194
195#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
197#[serde(rename_all = "snake_case")]
198pub enum FinishReason {
199 Stop,
201 Length,
203 ContentFilter,
205 ToolCalls,
207}
208
209impl FinishReason {
210 pub fn as_str(self) -> &'static str {
212 match self {
213 Self::Stop => "stop",
214 Self::Length => "length",
215 Self::ContentFilter => "content_filter",
216 Self::ToolCalls => "tool_calls",
217 }
218 }
219}
220
221#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
223pub struct Usage {
224 pub prompt_tokens: u32,
226 pub completion_tokens: u32,
228 pub total_tokens: u32,
230 #[serde(
232 skip_serializing_if = "Option::is_none",
233 default,
234 alias = "thinking_tokens"
235 )]
236 pub reasoning_tokens: Option<u32>,
237}
238
239impl Usage {
240 pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
242 Self {
243 prompt_tokens,
244 completion_tokens,
245 total_tokens: prompt_tokens + completion_tokens,
246 reasoning_tokens: None,
247 }
248 }
249}
250
251#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
253pub struct CompletionChunk {
254 pub id: String,
256 pub model: String,
258 pub choices: Vec<ChoiceDelta>,
260 #[serde(skip_serializing_if = "Option::is_none")]
262 pub created: Option<i64>,
263 #[serde(skip_serializing_if = "Option::is_none")]
265 pub usage: Option<Usage>,
266}
267
268#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
270pub struct ChoiceDelta {
271 pub index: u32,
273 pub delta: MessageDelta,
275 #[serde(skip_serializing_if = "Option::is_none")]
277 pub finish_reason: Option<FinishReason>,
278}
279
280#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
282pub struct MessageDelta {
283 #[serde(skip_serializing_if = "Option::is_none")]
285 pub role: Option<crate::message::Role>,
286 #[serde(skip_serializing_if = "Option::is_none")]
288 pub content: Option<String>,
289 #[serde(skip_serializing_if = "Option::is_none")]
291 pub reasoning_content: Option<String>,
292 #[serde(skip_serializing_if = "Option::is_none")]
294 pub tool_calls: Option<Vec<ToolCallDelta>>,
295}
296
297#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
299pub struct ToolCallDelta {
300 pub index: u32,
302 #[serde(skip_serializing_if = "Option::is_none")]
304 pub id: Option<String>,
305 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
307 pub tool_type: Option<crate::tool::ToolType>,
308 #[serde(skip_serializing_if = "Option::is_none")]
310 pub function: Option<ToolCallFunctionDelta>,
311}
312
313#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
315pub struct ToolCallFunctionDelta {
316 #[serde(skip_serializing_if = "Option::is_none")]
318 pub name: Option<String>,
319 #[serde(skip_serializing_if = "Option::is_none")]
321 pub arguments: Option<String>,
322}
323
324#[cfg(test)]
325mod tests {
326 use super::*;
327
328 #[test]
329 fn test_completion_response_content() {
330 let response = CompletionResponse {
331 id: "resp_123".to_string(),
332 model: "gpt-4".to_string(),
333 choices: vec![CompletionChoice {
334 index: 0,
335 message: Message::assistant("Hello!"),
336 finish_reason: FinishReason::Stop,
337 logprobs: None,
338 }],
339 usage: Usage::new(10, 5),
340 created: Some(1234567890),
341 provider: Some("openai".to_string()),
342 healing_metadata: None,
343 };
344
345 assert_eq!(response.content(), Some("Hello!"));
346 assert_eq!(response.first_choice().unwrap().index, 0);
347 assert!(!response.was_healed());
348 assert_eq!(response.confidence(), 1.0);
349 }
350
351 #[test]
352 fn test_completion_response_empty_choices() {
353 let response = CompletionResponse {
354 id: "resp_123".to_string(),
355 model: "gpt-4".to_string(),
356 choices: vec![],
357 usage: Usage::new(10, 0),
358 created: None,
359 provider: None,
360 healing_metadata: None,
361 };
362
363 assert_eq!(response.content(), None);
364 assert_eq!(response.first_choice(), None);
365 }
366
367 #[test]
368 fn test_usage_calculation() {
369 let usage = Usage::new(100, 50);
370 assert_eq!(usage.prompt_tokens, 100);
371 assert_eq!(usage.completion_tokens, 50);
372 assert_eq!(usage.total_tokens, 150);
373 assert_eq!(usage.reasoning_tokens, None);
374 }
375
376 #[test]
377 fn test_usage_deserializes_thinking_tokens_alias() {
378 let json = serde_json::json!({
379 "prompt_tokens": 10,
380 "completion_tokens": 5,
381 "total_tokens": 15,
382 "thinking_tokens": 3
383 });
384
385 let usage: Usage = serde_json::from_value(json).unwrap();
386 assert_eq!(usage.reasoning_tokens, Some(3));
387 }
388
389 #[test]
390 fn test_usage_serializes_reasoning_tokens_name() {
391 let usage = Usage {
392 prompt_tokens: 10,
393 completion_tokens: 5,
394 total_tokens: 15,
395 reasoning_tokens: Some(3),
396 };
397
398 let json = serde_json::to_value(&usage).unwrap();
399 assert_eq!(
400 json.get("reasoning_tokens").and_then(|v| v.as_u64()),
401 Some(3)
402 );
403 assert!(json.get("thinking_tokens").is_none());
404 }
405
406 #[test]
407 fn test_finish_reason_serialization() {
408 let json = serde_json::to_string(&FinishReason::Stop).unwrap();
409 assert_eq!(json, "\"stop\"");
410
411 let json = serde_json::to_string(&FinishReason::Length).unwrap();
412 assert_eq!(json, "\"length\"");
413
414 let json = serde_json::to_string(&FinishReason::ContentFilter).unwrap();
415 assert_eq!(json, "\"content_filter\"");
416
417 let json = serde_json::to_string(&FinishReason::ToolCalls).unwrap();
418 assert_eq!(json, "\"tool_calls\"");
419 }
420
421 #[test]
422 fn test_response_serialization() {
423 let response = CompletionResponse {
424 id: "resp_123".to_string(),
425 model: "gpt-4".to_string(),
426 choices: vec![CompletionChoice {
427 index: 0,
428 message: Message::assistant("Hello!"),
429 finish_reason: FinishReason::Stop,
430 logprobs: None,
431 }],
432 usage: Usage::new(10, 5),
433 created: None,
434 provider: None,
435 healing_metadata: None,
436 };
437
438 let json = serde_json::to_string(&response).unwrap();
439 let parsed: CompletionResponse = serde_json::from_str(&json).unwrap();
440 assert_eq!(response, parsed);
441 }
442
443 #[test]
444 fn test_streaming_chunk() {
445 let chunk = CompletionChunk {
446 id: "resp_123".to_string(),
447 model: "gpt-4".to_string(),
448 choices: vec![ChoiceDelta {
449 index: 0,
450 delta: MessageDelta {
451 role: Some(crate::message::Role::Assistant),
452 content: Some("Hello".to_string()),
453 reasoning_content: None,
454 tool_calls: None,
455 },
456 finish_reason: None,
457 }],
458 created: Some(1234567890),
459 usage: None,
460 };
461
462 let json = serde_json::to_string(&chunk).unwrap();
463 let parsed: CompletionChunk = serde_json::from_str(&json).unwrap();
464 assert_eq!(chunk, parsed);
465 }
466
467 #[test]
468 fn test_message_delta() {
469 let delta = MessageDelta {
470 role: Some(crate::message::Role::Assistant),
471 content: Some("Hi".to_string()),
472 reasoning_content: None,
473 tool_calls: None,
474 };
475
476 let json = serde_json::to_value(&delta).unwrap();
477 assert_eq!(json.get("role").and_then(|v| v.as_str()), Some("assistant"));
478 assert_eq!(json.get("content").and_then(|v| v.as_str()), Some("Hi"));
479 }
480
481 #[test]
482 fn test_optional_fields_not_serialized() {
483 let response = CompletionResponse {
484 id: "resp_123".to_string(),
485 model: "gpt-4".to_string(),
486 choices: vec![],
487 usage: Usage::new(10, 5),
488 created: None,
489 provider: None,
490 healing_metadata: None,
491 };
492
493 let json = serde_json::to_value(&response).unwrap();
494 assert!(json.get("created").is_none());
495 assert!(json.get("provider").is_none());
496 assert!(json.get("healing_metadata").is_none());
497 }
498
499 #[test]
500 fn test_healing_metadata() {
501 use crate::coercion::CoercionFlag;
502
503 let metadata = HealingMetadata::new(
504 vec![CoercionFlag::StrippedMarkdown],
505 0.9,
506 "Parse error".to_string(),
507 );
508
509 assert_eq!(metadata.confidence, 0.9);
510 assert!(!metadata.has_major_coercions());
511 assert!(metadata.is_confident(0.8));
512 assert!(!metadata.is_confident(0.95));
513
514 let major_metadata = HealingMetadata::new(
515 vec![CoercionFlag::TruncatedJson],
516 0.7,
517 "Parse error".to_string(),
518 );
519
520 assert!(major_metadata.has_major_coercions());
521 }
522
523 #[test]
524 fn test_healing_metadata_confidence_clamped() {
525 let metadata = HealingMetadata::new(vec![], 1.5, "error".to_string());
526 assert_eq!(metadata.confidence, 1.0);
527
528 let metadata = HealingMetadata::new(vec![], -0.5, "error".to_string());
529 assert_eq!(metadata.confidence, 0.0);
530 }
531
532 #[test]
533 fn test_response_with_healing_metadata() {
534 use crate::coercion::CoercionFlag;
535
536 let metadata = HealingMetadata::new(
537 vec![
538 CoercionFlag::StrippedMarkdown,
539 CoercionFlag::FixedTrailingComma,
540 ],
541 0.85,
542 "JSON parse error".to_string(),
543 );
544
545 let response = CompletionResponse {
546 id: "resp_123".to_string(),
547 model: "gpt-4".to_string(),
548 choices: vec![],
549 usage: Usage::new(10, 5),
550 created: None,
551 provider: None,
552 healing_metadata: Some(metadata),
553 };
554
555 assert!(response.was_healed());
556 assert_eq!(response.confidence(), 0.85);
557
558 let json = serde_json::to_string(&response).unwrap();
559 let parsed: CompletionResponse = serde_json::from_str(&json).unwrap();
560 assert_eq!(response, parsed);
561 assert!(parsed.was_healed());
562 assert_eq!(parsed.confidence(), 0.85);
563 }
564}