batch_mode_process_response/
handle_successful_response.rs

1// ---------------- [ File: batch-mode-process-response/src/handle_successful_response.rs ]
2crate::ix!();
3
4#[instrument(level="trace", skip_all)]
5pub async fn handle_successful_response<T>(
6    success_body:          &BatchSuccessResponseBody,
7    workspace:             &dyn BatchWorkspaceInterface,
8    expected_content_type: &ExpectedContentType,
9) -> Result<(), BatchSuccessResponseHandlingError> 
10where
11    T: 'static + Send + Sync + Named + DeserializeOwned + GetTargetPathForAIExpansion,
12{
13    trace!("Entering handle_successful_response with success_body ID: {}", success_body.id());
14    trace!("success_body => finish_reason={:?}, total_choices={}",
15        success_body.choices().get(0).map(|c| c.finish_reason()),
16        success_body.choices().len()
17    );
18
19    let choice = &success_body.choices()[0];
20    let message_content = choice.message().content();
21    trace!("Pulled first choice => finish_reason={:?}", choice.finish_reason());
22
23    if *choice.finish_reason() == FinishReason::Length {
24        trace!("Detected finish_reason=Length => calling handle_finish_reason_length");
25        handle_finish_reason_length(success_body.id(), message_content).await?;
26        trace!("Returned from handle_finish_reason_length with success_body ID: {}", success_body.id());
27    }
28
29    match expected_content_type {
30        ExpectedContentType::Json => {
31            trace!("ExpectedContentType::Json => about to extract/repair JSON for success_body ID: {}", success_body.id());
32            match message_content.extract_clean_parse_json_with_repair() {
33                Ok(json_content) => {
34                    debug!("JSON parse/repair succeeded for success_body ID: {}", success_body.id());
35                    trace!("Now deserializing into typed struct T...");
36
37                    // In handle_successful_response.rs:
38                    let typed_item: T = match serde_json::from_value(json_content.clone()) {
39                        Ok(t) => {
40                            trace!("Deserialization into T succeeded...");
41                            t
42                        }
43                        Err(e) => {
44                            error!("Deserialization into T failed: {:?}", e);
45                            // We also call handle_failed_json_repair here so that the test sees a file
46                            handle_failed_json_repair(success_body.id(), message_content, workspace).await?;
47                            return Err(e.into());
48                        }
49                    };
50
51                    // Convert to Arc if needed
52                    trace!("Wrapping typed_item in Arc => T::name()={}", typed_item.name());
53                    let typed_item_arc: Arc<dyn GetTargetPathForAIExpansion + Send + Sync + 'static> = Arc::new(typed_item);
54
55                    // Determine target path
56                    let target_path = workspace.target_path(&typed_item_arc, expected_content_type);
57                    trace!("Target path computed => {:?}", target_path);
58
59                    // Pretty-print the fully repaired JSON
60                    let serialized_json = match serde_json::to_string_pretty(&json_content) {
61                        Ok(s) => {
62                            trace!("Successfully created pretty JSON string for success_body ID: {}", success_body.id());
63                            s
64                        }
65                        Err(e) => {
66                            error!("Re-serialization to pretty JSON failed: {:?}", e);
67                            return Err(JsonParseError::SerdeError(e).into());
68                        }
69                    };
70
71                    info!("writing JSON output to {:?}", target_path);
72                    write_to_file(&target_path, &serialized_json).await?;
73                    trace!("Successfully wrote JSON file => {:?}", target_path);
74                    trace!("Exiting handle_successful_response with success_body ID: {}", success_body.id());
75                    Ok(())
76                }
77                Err(e) => {
78                    warn!("JSON extraction/repair failed for success_body ID: {} with error: {:?}", success_body.id(), e);
79                    let failed_id = success_body.id();
80                    trace!("Calling handle_failed_json_repair for ID={}", failed_id);
81                    handle_failed_json_repair(failed_id, message_content, workspace).await?;
82                    trace!("Returned from handle_failed_json_repair => now returning error for ID={}", failed_id);
83                    Err(e.into())
84                }
85            }
86        }
87        ExpectedContentType::PlainText => {
88            trace!("Received plain text content for request {} => length={}", success_body.id(), message_content.len());
89            let index = BatchIndex::from_uuid_str(success_body.id())?;
90            trace!("Parsed BatchIndex => {:?}", index);
91
92            let text_path = workspace.text_storage_path(&index);
93            info!("writing plain text output to {:?}", text_path);
94            write_to_file(&text_path, message_content.as_str()).await?;
95            trace!("Successfully wrote plain text file => {:?}", text_path);
96
97            trace!("Exiting handle_successful_response with success_body ID: {}", success_body.id());
98            Ok(())
99        }
100        _ => { todo!() }
101    }
102}
103
104#[cfg(test)]
105mod handle_successful_response_tests {
106    use super::*;
107    use std::fs;
108
109    #[derive(Debug, Deserialize, Serialize, NamedItem)]
110    pub struct MockItemForSuccess {
111        pub name: String,
112    }
113
114    #[traced_test]
115    async fn test_handle_successful_response_json_failure() {
116        // This test tries to parse invalid JSON into our `MockItemForSuccess`,
117        // expecting it to fail and log a file in `failed_json_repairs_dir`.
118        trace!("===== BEGIN TEST: test_handle_successful_response_json_failure =====");
119
120        // 1) Use ephemeral default workspace (no overrides).
121        let workspace = BatchWorkspace::new_temp().await.unwrap();
122        info!("Created ephemeral workspace: {:?}", workspace);
123
124        // 2) Ensure repairs dir is empty
125        let repairs_dir = workspace.failed_json_repairs_dir();
126
127        // 3) Create a response that is *not* valid JSON
128        let invalid_msg = ChatCompletionResponseMessage {
129            role: Role::Assistant,
130            content: Some("this is not valid json at all".into()),
131            audio: None,
132            function_call: None,
133            refusal: None,
134            tool_calls: None,
135        };
136
137        let choice_fail = BatchChoiceBuilder::default()
138            .index(0_u32)
139            .finish_reason(FinishReason::Stop)
140            .logprobs(None)
141            .message(invalid_msg)
142            .build()
143            .unwrap();
144
145        let success_body = BatchSuccessResponseBodyBuilder::default()
146            .object("response".to_string())
147            .id("some-other-uuid".to_string())
148            .created(0_u64)
149            .model("test-model".to_string())
150            .choices(vec![choice_fail])
151            .usage(BatchUsage::mock())
152            .build()
153            .unwrap();
154
155        // 4) Call handle_successful_response with ExpectedContentType=Json
156        let rc = handle_successful_response::<MockItemForSuccess>(
157            &success_body,
158            workspace.as_ref(),
159            &ExpectedContentType::Json
160        ).await;
161
162        // 5) Confirm it fails
163        assert!(rc.is_err(), "We expect an error due to invalid JSON content");
164
165        // 6) Confirm the "some-other-uuid" file is in the ephemeral repairs dir
166        let repair_path = repairs_dir.join("some-other-uuid");
167        trace!("Asserting that repair file path exists: {:?}", repair_path);
168        assert!(repair_path.exists(), "A repair file must be created for invalid JSON");
169
170        trace!("===== END TEST: test_handle_successful_response_json_failure =====");
171    }
172}