gaze_proxy/adapters/
gemini.rs

1use http::Method;
2use serde_json::Value;
3use url::Url;
4
5use crate::adapter::{walk_all_strings, PiiSurface, ProviderAdapter, SseEvent};
6
7#[derive(Clone, Debug)]
8#[non_exhaustive]
9pub struct GeminiAdapter {
10    upstream: Url,
11}
12
13impl GeminiAdapter {
14    pub fn new(upstream: Url) -> Self {
15        Self { upstream }
16    }
17}
18
19#[async_trait::async_trait]
20impl ProviderAdapter for GeminiAdapter {
21    fn name(&self) -> &'static str {
22        "gemini"
23    }
24
25    fn matches_path(&self, method: &Method, path: &str) -> bool {
26        method == Method::POST
27            && path.starts_with("/v1beta/models/")
28            && (path.ends_with(":generateContent")
29                || path.ends_with(":streamGenerateContent")
30                || path.ends_with(":countTokens"))
31    }
32
33    fn upstream_base(&self) -> &Url {
34        &self.upstream
35    }
36
37    fn request_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
38        collect_gemini_surfaces(body, true)
39    }
40
41    fn response_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
42        collect_gemini_surfaces(body, false)
43    }
44
45    fn sse_event_pii_surfaces<'a>(&self, event: &'a mut SseEvent) -> Vec<PiiSurface<'a>> {
46        self.response_pii_surfaces(&mut event.data)
47    }
48}
49
50fn collect_gemini_surfaces(body: &mut Value, request: bool) -> Vec<PiiSurface<'_>> {
51    let mut surfaces = Vec::new();
52    if let Value::Object(root) = body {
53        if request {
54            for (key, value) in root {
55                match key.as_str() {
56                    "contents" => {
57                        if let Value::Array(contents) = value {
58                            for (index, content) in contents.iter_mut().enumerate() {
59                                collect_content(
60                                    &mut surfaces,
61                                    format!("contents[{index}]"),
62                                    content,
63                                );
64                            }
65                        }
66                    }
67                    "systemInstruction" => {
68                        collect_content(&mut surfaces, "systemInstruction".to_string(), value);
69                    }
70                    _ => {}
71                }
72            }
73        } else if let Some(Value::Array(candidates)) = root.get_mut("candidates") {
74            for (index, candidate) in candidates.iter_mut().enumerate() {
75                if let Value::Object(candidate) = candidate {
76                    if let Some(content) = candidate.get_mut("content") {
77                        collect_content(
78                            &mut surfaces,
79                            format!("candidates[{index}].content"),
80                            content,
81                        );
82                    }
83                }
84            }
85        }
86    }
87    surfaces
88}
89
90fn collect_content<'a>(surfaces: &mut Vec<PiiSurface<'a>>, prefix: String, value: &'a mut Value) {
91    let Value::Object(content) = value else {
92        return;
93    };
94    if let Some(Value::Array(parts)) = content.get_mut("parts") {
95        for (index, part) in parts.iter_mut().enumerate() {
96            let Value::Object(part) = part else {
97                continue;
98            };
99            for (key, value) in part {
100                match key.as_str() {
101                    "text" => {
102                        if let Value::String(text) = value {
103                            surfaces.push(PiiSurface {
104                                field_path: format!("{prefix}.parts[{index}].text"),
105                                text,
106                            });
107                        }
108                    }
109                    "functionCall" => {
110                        if let Some(args) =
111                            value.as_object_mut().and_then(|call| call.get_mut("args"))
112                        {
113                            walk_all_strings(
114                                surfaces,
115                                format!("{prefix}.parts[{index}].functionCall.args"),
116                                args,
117                            );
118                        }
119                    }
120                    "functionResponse" => {
121                        if let Some(response) = value
122                            .as_object_mut()
123                            .and_then(|call| call.get_mut("response"))
124                        {
125                            walk_all_strings(
126                                surfaces,
127                                format!("{prefix}.parts[{index}].functionResponse.response"),
128                                response,
129                            );
130                        }
131                    }
132                    _ => {}
133                }
134            }
135        }
136    }
137}
gaze_proxy/adapters/gemini.rs

gaze_proxy/adapters/
gemini.rs
