Skip to main content

gproxy_protocol/transform/openai/websocket/from_http/
request.rs

1use std::collections::BTreeMap;
2
3use crate::openai::create_response::request::{OpenAiCreateResponseRequest, RequestBody};
4use crate::openai::create_response::types::Metadata;
5use crate::openai::create_response::websocket::request::OpenAiCreateResponseWebSocketConnectRequest;
6use crate::openai::create_response::websocket::types::{
7    OpenAiCreateResponseCreateWebSocketRequestBody, OpenAiCreateResponseWebSocketClientMessage,
8};
9use crate::transform::openai::websocket::context::OpenAiWebsocketTransformContext;
10use crate::transform::utils::TransformError;
11
12pub const OPENAI_CLIENT_METADATA_TUNNEL_PREFIX: &str = "gproxy.ws.client_metadata.";
13
14fn extract_client_metadata_from_request_body(
15    body: &mut RequestBody,
16    ctx: &mut OpenAiWebsocketTransformContext,
17) -> Option<Metadata> {
18    let metadata = body.metadata.as_mut()?;
19
20    let mut tunnel = BTreeMap::new();
21    let keys = metadata.keys().cloned().collect::<Vec<_>>();
22    for key in keys {
23        if let Some(tunnel_key) = key.strip_prefix(OPENAI_CLIENT_METADATA_TUNNEL_PREFIX) {
24            let Some(value) = metadata.remove(&key) else {
25                continue;
26            };
27            tunnel.insert(tunnel_key.to_string(), value);
28        }
29    }
30
31    if metadata.is_empty() {
32        body.metadata = None;
33    }
34
35    if !tunnel.is_empty() {
36        ctx.push_warning(
37            "openai websocket from_http request: restored client_metadata from metadata tunnel",
38        );
39        Some(tunnel)
40    } else {
41        None
42    }
43}
44
45pub fn openai_create_response_request_to_websocket_message_with_context(
46    value: &OpenAiCreateResponseRequest,
47) -> Result<
48    (
49        OpenAiCreateResponseWebSocketClientMessage,
50        OpenAiWebsocketTransformContext,
51    ),
52    TransformError,
53> {
54    let mut ctx = OpenAiWebsocketTransformContext::default();
55    let mut request_body = value.body.clone();
56    let client_metadata = extract_client_metadata_from_request_body(&mut request_body, &mut ctx);
57
58    Ok((
59        OpenAiCreateResponseWebSocketClientMessage::ResponseCreate(
60            OpenAiCreateResponseCreateWebSocketRequestBody {
61                request: request_body,
62                generate: None,
63                client_metadata,
64            },
65        ),
66        ctx,
67    ))
68}
69
70pub fn openai_create_response_request_to_websocket_connect_with_context(
71    value: &OpenAiCreateResponseRequest,
72) -> Result<
73    (
74        OpenAiCreateResponseWebSocketConnectRequest,
75        OpenAiWebsocketTransformContext,
76    ),
77    TransformError,
78> {
79    let (message, ctx) = openai_create_response_request_to_websocket_message_with_context(value)?;
80    Ok((
81        OpenAiCreateResponseWebSocketConnectRequest {
82            body: Some(message),
83            ..OpenAiCreateResponseWebSocketConnectRequest::default()
84        },
85        ctx,
86    ))
87}
88
89impl TryFrom<&OpenAiCreateResponseRequest> for OpenAiCreateResponseWebSocketClientMessage {
90    type Error = TransformError;
91
92    fn try_from(value: &OpenAiCreateResponseRequest) -> Result<Self, TransformError> {
93        Ok(openai_create_response_request_to_websocket_message_with_context(value)?.0)
94    }
95}
96
97impl TryFrom<&OpenAiCreateResponseRequest> for OpenAiCreateResponseWebSocketConnectRequest {
98    type Error = TransformError;
99
100    fn try_from(value: &OpenAiCreateResponseRequest) -> Result<Self, TransformError> {
101        Ok(openai_create_response_request_to_websocket_connect_with_context(value)?.0)
102    }
103}
104
105impl TryFrom<OpenAiCreateResponseRequest> for OpenAiCreateResponseWebSocketConnectRequest {
106    type Error = TransformError;
107
108    fn try_from(value: OpenAiCreateResponseRequest) -> Result<Self, TransformError> {
109        OpenAiCreateResponseWebSocketConnectRequest::try_from(&value)
110    }
111}