Skip to main content

simple_agent_type/
response.rs

1//! Response types for LLM completions.
2//!
3//! Provides OpenAI-compatible response structures.
4
5use crate::coercion::CoercionFlag;
6use crate::message::Message;
7use serde::{Deserialize, Serialize};
8
9/// Metadata about healing/coercion applied to a response.
10///
11/// When native structured output parsing fails and healing is enabled,
12/// this metadata tracks all transformations applied to recover the response.
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct HealingMetadata {
15    /// All coercion flags applied during healing
16    pub flags: Vec<CoercionFlag>,
17    /// Confidence score (0.0-1.0) of the healed response
18    pub confidence: f32,
19    /// The original parsing error that triggered healing
20    pub original_error: String,
21}
22
23impl HealingMetadata {
24    /// Create new healing metadata.
25    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    /// Check if any major coercions were applied.
34    pub fn has_major_coercions(&self) -> bool {
35        self.flags.iter().any(|f| f.is_major())
36    }
37
38    /// Check if confidence meets a threshold.
39    pub fn is_confident(&self, threshold: f32) -> bool {
40        self.confidence >= threshold
41    }
42}
43
44/// A completion response from an LLM provider.
45#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
46pub struct CompletionResponse {
47    /// Unique response identifier
48    pub id: String,
49    /// Model used for completion
50    pub model: String,
51    /// List of completion choices
52    pub choices: Vec<CompletionChoice>,
53    /// Token usage statistics
54    pub usage: Usage,
55    /// Unix timestamp of creation
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub created: Option<i64>,
58    /// Provider that generated this response
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub provider: Option<String>,
61    /// Healing metadata (present if response was healed after parse failure)
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub healing_metadata: Option<HealingMetadata>,
64}
65
66impl CompletionResponse {
67    /// Get the content of the first choice (convenience method).
68    ///
69    /// # Example
70    /// ```
71    /// use simple_agent_type::response::{CompletionResponse, CompletionChoice, Usage, FinishReason};
72    /// use simple_agent_type::message::Message;
73    ///
74    /// let response = CompletionResponse {
75    ///     id: "resp_123".to_string(),
76    ///     model: "gpt-4".to_string(),
77    ///     choices: vec![CompletionChoice {
78    ///         index: 0,
79    ///         message: Message::assistant("Hello!"),
80    ///         finish_reason: FinishReason::Stop,
81    ///         logprobs: None,
82    ///     }],
83    ///     usage: Usage {
84    ///         prompt_tokens: 10,
85    ///         completion_tokens: 5,
86    ///         total_tokens: 15,
87    ///         reasoning_tokens: None,
88    ///     },
89    ///     created: None,
90    ///     provider: None,
91    ///     healing_metadata: None,
92    /// };
93    ///
94    /// assert_eq!(response.content(), Some("Hello!"));
95    /// ```
96    pub fn content(&self) -> Option<&str> {
97        self.choices
98            .first()
99            .map(|choice| choice.message.content.as_str())
100    }
101
102    /// Get the first choice.
103    pub fn first_choice(&self) -> Option<&CompletionChoice> {
104        self.choices.first()
105    }
106
107    /// Check if this response was healed after a parsing failure.
108    ///
109    /// Returns `true` if healing metadata is present, indicating the response
110    /// required transformation to be parseable.
111    ///
112    /// # Example
113    /// ```
114    /// use simple_agent_type::response::{CompletionResponse, CompletionChoice, Usage, FinishReason, HealingMetadata};
115    /// use simple_agent_type::message::Message;
116    /// use simple_agent_type::coercion::CoercionFlag;
117    ///
118    /// let mut response = CompletionResponse {
119    ///     id: "resp_123".to_string(),
120    ///     model: "gpt-4".to_string(),
121    ///     choices: vec![],
122    ///     usage: Usage::new(10, 5),
123    ///     created: None,
124    ///     provider: None,
125    ///     healing_metadata: None,
126    /// };
127    ///
128    /// assert!(!response.was_healed());
129    ///
130    /// response.healing_metadata = Some(HealingMetadata::new(
131    ///     vec![CoercionFlag::StrippedMarkdown],
132    ///     0.9,
133    ///     "Parse error".to_string(),
134    /// ));
135    ///
136    /// assert!(response.was_healed());
137    /// ```
138    pub fn was_healed(&self) -> bool {
139        self.healing_metadata.is_some()
140    }
141
142    /// Get the confidence score of the response.
143    ///
144    /// Returns 1.0 if the response was not healed (perfect confidence),
145    /// otherwise returns the confidence score from healing metadata.
146    ///
147    /// # Example
148    /// ```
149    /// use simple_agent_type::response::{CompletionResponse, CompletionChoice, Usage, FinishReason, HealingMetadata};
150    /// use simple_agent_type::message::Message;
151    /// use simple_agent_type::coercion::CoercionFlag;
152    ///
153    /// let mut response = CompletionResponse {
154    ///     id: "resp_123".to_string(),
155    ///     model: "gpt-4".to_string(),
156    ///     choices: vec![],
157    ///     usage: Usage::new(10, 5),
158    ///     created: None,
159    ///     provider: None,
160    ///     healing_metadata: None,
161    /// };
162    ///
163    /// assert_eq!(response.confidence(), 1.0);
164    ///
165    /// response.healing_metadata = Some(HealingMetadata::new(
166    ///     vec![CoercionFlag::StrippedMarkdown],
167    ///     0.8,
168    ///     "Parse error".to_string(),
169    /// ));
170    ///
171    /// assert_eq!(response.confidence(), 0.8);
172    /// ```
173    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/// A single completion choice.
182#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
183pub struct CompletionChoice {
184    /// Index of this choice
185    pub index: u32,
186    /// The message content
187    pub message: Message,
188    /// Why the completion finished
189    pub finish_reason: FinishReason,
190    /// Log probabilities (if requested)
191    #[serde(skip_serializing_if = "Option::is_none")]
192    pub logprobs: Option<serde_json::Value>,
193}
194
195/// Reason why a completion finished.
196#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
197#[serde(rename_all = "snake_case")]
198pub enum FinishReason {
199    /// Natural stop point reached
200    Stop,
201    /// Maximum token length reached
202    Length,
203    /// Content filtered by provider
204    ContentFilter,
205    /// Tool/function calls generated
206    ToolCalls,
207}
208
209impl FinishReason {
210    /// Returns this finish reason as its canonical snake_case string value.
211    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/// Token usage statistics.
222#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
223pub struct Usage {
224    /// Tokens in the prompt
225    pub prompt_tokens: u32,
226    /// Tokens in the completion
227    pub completion_tokens: u32,
228    /// Total tokens used
229    pub total_tokens: u32,
230    /// Optional reasoning token usage (provider-dependent).
231    #[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    /// Create a new Usage with calculated total.
241    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/// A chunk of a streaming completion response.
252#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
253pub struct CompletionChunk {
254    /// Unique response identifier
255    pub id: String,
256    /// Model used for completion
257    pub model: String,
258    /// List of choice deltas
259    pub choices: Vec<ChoiceDelta>,
260    /// Unix timestamp of creation
261    #[serde(skip_serializing_if = "Option::is_none")]
262    pub created: Option<i64>,
263    /// Optional token usage for this chunk (typically on final chunk)
264    #[serde(skip_serializing_if = "Option::is_none")]
265    pub usage: Option<Usage>,
266}
267
268/// A delta in a streaming choice.
269#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
270pub struct ChoiceDelta {
271    /// Index of this choice
272    pub index: u32,
273    /// The message delta
274    pub delta: MessageDelta,
275    /// Why the completion finished (only in final chunk)
276    #[serde(skip_serializing_if = "Option::is_none")]
277    pub finish_reason: Option<FinishReason>,
278}
279
280/// Incremental message content in a stream.
281#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
282pub struct MessageDelta {
283    /// Role (only in first chunk)
284    #[serde(skip_serializing_if = "Option::is_none")]
285    pub role: Option<crate::message::Role>,
286    /// Incremental content
287    #[serde(skip_serializing_if = "Option::is_none")]
288    pub content: Option<String>,
289    /// Optional incremental reasoning/thinking content.
290    #[serde(skip_serializing_if = "Option::is_none")]
291    pub reasoning_content: Option<String>,
292    /// Optional incremental tool call deltas.
293    #[serde(skip_serializing_if = "Option::is_none")]
294    pub tool_calls: Option<Vec<ToolCallDelta>>,
295}
296
297/// Incremental tool call payload emitted in streaming responses.
298#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
299pub struct ToolCallDelta {
300    /// Tool call position in the choice stream.
301    pub index: u32,
302    /// Tool call identifier (may arrive incrementally).
303    #[serde(skip_serializing_if = "Option::is_none")]
304    pub id: Option<String>,
305    /// Tool type.
306    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
307    pub tool_type: Option<crate::tool::ToolType>,
308    /// Function name/arguments payload.
309    #[serde(skip_serializing_if = "Option::is_none")]
310    pub function: Option<ToolCallFunctionDelta>,
311}
312
313/// Incremental function payload for a streamed tool call.
314#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
315pub struct ToolCallFunctionDelta {
316    /// Function name (may arrive once).
317    #[serde(skip_serializing_if = "Option::is_none")]
318    pub name: Option<String>,
319    /// JSON arguments text (may arrive in chunks).
320    #[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}