batch_mode_process_response/
handle_successful_response.rs1crate::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 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 handle_failed_json_repair(success_body.id(), message_content, workspace).await?;
47 return Err(e.into());
48 }
49 };
50
51 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 let target_path = workspace.target_path(&typed_item_arc, expected_content_type);
57 trace!("Target path computed => {:?}", target_path);
58
59 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 trace!("===== BEGIN TEST: test_handle_successful_response_json_failure =====");
119
120 let workspace = BatchWorkspace::new_temp().await.unwrap();
122 info!("Created ephemeral workspace: {:?}", workspace);
123
124 let repairs_dir = workspace.failed_json_repairs_dir();
126
127 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 let rc = handle_successful_response::<MockItemForSuccess>(
157 &success_body,
158 workspace.as_ref(),
159 &ExpectedContentType::Json
160 ).await;
161
162 assert!(rc.is_err(), "We expect an error due to invalid JSON content");
164
165 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}