gproxy-transform 3.0.0-alpha.0

GPROXY v3 pairwise protocol transforms
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use bytes::Bytes;
use gproxy_protocol::gemini;
use gproxy_protocol::openai::images as openai_images;

use crate::TransformError;

pub(crate) fn openai_request(body: Bytes) -> Result<Bytes, TransformError> {
    let input: openai_images::CreateImageRequest = serde_json::from_slice(&body)?;
    super::encode(&openai_request_typed(input))
}

pub(crate) fn openai_request_typed(
    input: openai_images::CreateImageRequest,
) -> gemini::ImagenPredictRequest {
    let parameters = serde_json::json!({
        "sampleCount": input.n.unwrap_or(1),
        "outputOptions": {
            "mimeType": input.output_format.as_ref().and_then(wire).map(|v| format!("image/{v}"))
        },
        "aspectRatio": input.size.as_ref().and_then(wire)
    });
    crate::wire!(gemini::ImagenPredictRequest {
        instances: vec![serde_json::json!({"prompt": input.prompt})],
        parameters: Some(parameters),
        rest: Default::default(),
    })
}

pub(crate) fn gemini_request(body: Bytes, model: &str) -> Result<Bytes, TransformError> {
    let input: gemini::ImagenPredictRequest = serde_json::from_slice(&body)?;
    super::encode(&gemini_request_typed(input, model))
}

pub(crate) fn gemini_request_typed(
    input: gemini::ImagenPredictRequest,
    model: &str,
) -> openai_images::CreateImageRequest {
    let prompt = input
        .instances
        .iter()
        .find_map(|instance| instance.get("prompt").and_then(serde_json::Value::as_str))
        .unwrap_or_default()
        .to_owned();
    let n = input
        .parameters
        .as_ref()
        .and_then(|value| value.get("sampleCount"))
        .and_then(serde_json::Value::as_u64)
        .and_then(|value| u32::try_from(value).ok());
    crate::wire!(openai_images::CreateImageRequest {
        prompt,
        model: Some(model.into()),
        n,
        rest: Default::default(),
        background: None,
        moderation: None,
        output_compression: None,
        output_format: None,
        partial_images: None,
        quality: None,
        response_format: None,
        size: None,
        stream: None,
        style: None,
        user: None,
    })
}

pub(crate) fn gemini_response_to_openai(body: Bytes) -> Result<Bytes, TransformError> {
    let input: gemini::ImagenPredictResponse = serde_json::from_slice(&body)?;
    super::encode(&gemini_response_to_openai_typed(input))
}

pub(crate) fn gemini_response_to_openai_typed(
    input: gemini::ImagenPredictResponse,
) -> openai_images::ImagesResponse {
    let data = input
        .predictions
        .into_iter()
        .filter_map(|prediction| {
            let b64_json = prediction
                .get("bytesBase64Encoded")
                .or_else(|| prediction.pointer("/image/bytesBase64Encoded"))
                .and_then(serde_json::Value::as_str)?
                .to_owned();
            Some(crate::wire!(openai_images::Image {
                b64_json: Some(b64_json),
                revised_prompt: prediction
                    .get("prompt")
                    .and_then(serde_json::Value::as_str)
                    .map(str::to_owned),
                url: None,
                rest: Default::default(),
            }))
        })
        .collect();
    crate::wire!(openai_images::ImagesResponse {
        created: 0,
        data: Some(data),
        rest: Default::default(),
        background: None,
        output_format: None,
        quality: None,
        size: None,
        usage: None,
    })
}

pub(crate) fn openai_response_to_gemini(body: Bytes) -> Result<Bytes, TransformError> {
    let input: openai_images::ImagesResponse = serde_json::from_slice(&body)?;
    super::encode(&openai_response_to_gemini_typed(input))
}

pub(crate) fn openai_response_to_gemini_typed(
    input: openai_images::ImagesResponse,
) -> gemini::ImagenPredictResponse {
    let predictions = input
        .data
        .unwrap_or_default()
        .into_iter()
        .filter_map(|image| {
            image.b64_json.map(|data| {
                serde_json::json!({
                    "bytesBase64Encoded": data,
                    "prompt": image.revised_prompt
                })
            })
        })
        .collect();
    crate::wire!(gemini::ImagenPredictResponse {
        predictions,
        rest: Default::default(),
    })
}

fn wire(value: &impl serde::Serialize) -> Option<String> {
    serde_json::to_value(value)
        .ok()?
        .as_str()
        .map(str::to_owned)
}