gaze_proxy/adapters/
openai.rs

1use http::Method;
2use serde_json::Value;
3use url::Url;
4
5use crate::adapter::{
6    push_string, push_text_blocks, walk_all_strings, PiiSurface, ProviderAdapter, SseEvent,
7};
8
9#[derive(Clone, Debug)]
10#[non_exhaustive]
11pub struct OpenAiAdapter {
12    upstream: Url,
13}
14
15impl OpenAiAdapter {
16    pub fn new(upstream: Url) -> Self {
17        Self { upstream }
18    }
19}
20
21#[async_trait::async_trait]
22impl ProviderAdapter for OpenAiAdapter {
23    fn name(&self) -> &'static str {
24        "openai"
25    }
26
27    fn matches_path(&self, method: &Method, path: &str) -> bool {
28        method == Method::POST
29            && matches!(
30                path,
31                "/v1/chat/completions" | "/v1/completions" | "/v1/responses"
32            )
33    }
34
35    fn upstream_base(&self) -> &Url {
36        &self.upstream
37    }
38
39    fn request_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
40        let mut surfaces = Vec::new();
41        if let Value::Object(root) = body {
42            for (key, value) in root {
43                match key.as_str() {
44                    "system" => push_string(&mut surfaces, "system", value),
45                    "prompt" => push_text_blocks(&mut surfaces, "prompt", value),
46                    "instructions" => push_string(&mut surfaces, "instructions", value),
47                    "messages" => {
48                        if let Value::Array(messages) = value {
49                            for (index, message) in messages.iter_mut().enumerate() {
50                                collect_message_surfaces(
51                                    &mut surfaces,
52                                    format!("messages[{index}]"),
53                                    message,
54                                );
55                            }
56                        }
57                    }
58                    "input" => push_text_blocks(&mut surfaces, "input", value),
59                    _ => {}
60                }
61            }
62        }
63        surfaces
64    }
65
66    fn response_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
67        let mut surfaces = Vec::new();
68        if let Value::Object(root) = body {
69            for (key, value) in root {
70                match key.as_str() {
71                    "choices" => {
72                        if let Value::Array(choices) = value {
73                            for (index, choice) in choices.iter_mut().enumerate() {
74                                if let Value::Object(choice) = choice {
75                                    for (choice_key, choice_value) in choice {
76                                        match choice_key.as_str() {
77                                            "message" => collect_message_surfaces(
78                                                &mut surfaces,
79                                                format!("choices[{index}].message"),
80                                                choice_value,
81                                            ),
82                                            "delta" => collect_message_surfaces(
83                                                &mut surfaces,
84                                                format!("choices[{index}].delta"),
85                                                choice_value,
86                                            ),
87                                            "text" => push_string(
88                                                &mut surfaces,
89                                                format!("choices[{index}].text"),
90                                                choice_value,
91                                            ),
92                                            _ => {}
93                                        }
94                                    }
95                                }
96                            }
97                        }
98                    }
99                    "output" => walk_all_strings(&mut surfaces, "output".to_string(), value),
100                    _ => {}
101                }
102            }
103        }
104        surfaces
105    }
106
107    fn sse_event_pii_surfaces<'a>(&self, event: &'a mut SseEvent) -> Vec<PiiSurface<'a>> {
108        self.response_pii_surfaces(&mut event.data)
109    }
110}
111
112fn collect_message_surfaces<'a>(
113    surfaces: &mut Vec<PiiSurface<'a>>,
114    prefix: String,
115    message: &'a mut Value,
116) {
117    let Value::Object(map) = message else {
118        return;
119    };
120    for (key, value) in map {
121        match key.as_str() {
122            "content" => push_text_blocks(surfaces, &format!("{prefix}.content"), value),
123            "tool_results" => push_text_blocks(surfaces, &format!("{prefix}.tool_results"), value),
124            "tool_calls" => {
125                if let Value::Array(tool_calls) = value {
126                    for (index, tool_call) in tool_calls.iter_mut().enumerate() {
127                        if let Value::Object(tool_call) = tool_call {
128                            if let Some(Value::Object(function)) = tool_call.get_mut("function") {
129                                if let Some(args) = function.get_mut("arguments") {
130                                    push_string(
131                                        surfaces,
132                                        format!("{prefix}.tool_calls[{index}].function.arguments"),
133                                        args,
134                                    );
135                                }
136                            }
137                        }
138                    }
139                }
140            }
141            _ => {}
142        }
143    }
144}
gaze_proxy/adapters/openai.rs

gaze_proxy/adapters/
openai.rs
