Skip to main content

gproxy_protocol/transform/openai/nonstream_to_stream/
response.rs

1use crate::openai::create_response::response::{OpenAiCreateResponseResponse, ResponseBody};
2use crate::openai::create_response::stream::ResponseStreamEvent;
3use crate::openai::create_response::types as rt;
4use crate::transform::utils::TransformError;
5
6fn with_status(mut body: ResponseBody, status: rt::ResponseStatus) -> ResponseBody {
7    body.status = Some(status);
8    body
9}
10
11fn take_sequence(next_sequence_number: &mut u64) -> u64 {
12    let sequence_number = *next_sequence_number;
13    *next_sequence_number = next_sequence_number.saturating_add(1);
14    sequence_number
15}
16
17impl TryFrom<OpenAiCreateResponseResponse> for Vec<ResponseStreamEvent> {
18    type Error = TransformError;
19
20    fn try_from(value: OpenAiCreateResponseResponse) -> Result<Self, TransformError> {
21        match value {
22            OpenAiCreateResponseResponse::Success { body, .. } => {
23                let mut events = Vec::new();
24                let mut next_sequence_number = 0_u64;
25
26                let in_progress = with_status(body.clone(), rt::ResponseStatus::InProgress);
27                let sequence_number = take_sequence(&mut next_sequence_number);
28                events.push(ResponseStreamEvent::Created {
29                    response: in_progress.clone(),
30                    sequence_number,
31                });
32                let sequence_number = take_sequence(&mut next_sequence_number);
33                events.push(ResponseStreamEvent::InProgress {
34                    response: in_progress,
35                    sequence_number,
36                });
37
38                let final_status = body.status.clone().unwrap_or_else(|| {
39                    if body.error.is_some() {
40                        rt::ResponseStatus::Failed
41                    } else if body.incomplete_details.is_some() {
42                        rt::ResponseStatus::Incomplete
43                    } else {
44                        rt::ResponseStatus::Completed
45                    }
46                });
47
48                match final_status {
49                    rt::ResponseStatus::Failed => {
50                        let sequence_number = take_sequence(&mut next_sequence_number);
51                        events.push(ResponseStreamEvent::Failed {
52                            response: with_status(body, rt::ResponseStatus::Failed),
53                            sequence_number,
54                        });
55                    }
56                    rt::ResponseStatus::Incomplete => {
57                        let sequence_number = take_sequence(&mut next_sequence_number);
58                        events.push(ResponseStreamEvent::Incomplete {
59                            response: with_status(body, rt::ResponseStatus::Incomplete),
60                            sequence_number,
61                        });
62                    }
63                    _ => {
64                        let sequence_number = take_sequence(&mut next_sequence_number);
65                        events.push(ResponseStreamEvent::Completed {
66                            response: with_status(body, rt::ResponseStatus::Completed),
67                            sequence_number,
68                        });
69                    }
70                }
71
72                Ok(events)
73            }
74            OpenAiCreateResponseResponse::Error { body, .. } => {
75                let error = body.error;
76                let events = vec![ResponseStreamEvent::Error {
77                    error: crate::openai::create_response::stream::ResponseStreamErrorPayload {
78                        type_: error.type_,
79                        code: error.code,
80                        message: error.message,
81                        param: error.param,
82                    },
83                    sequence_number: 0,
84                }];
85
86                Ok(events)
87            }
88        }
89    }
90}