Skip to main content

gaze_proxy/adapters/
anthropic.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 AnthropicAdapter {
12    upstream: Url,
13}
14
15impl AnthropicAdapter {
16    pub fn new(upstream: Url) -> Self {
17        Self { upstream }
18    }
19}
20
21#[async_trait::async_trait]
22impl ProviderAdapter for AnthropicAdapter {
23    fn name(&self) -> &'static str {
24        "anthropic"
25    }
26
27    fn matches_path(&self, method: &Method, path: &str) -> bool {
28        method == Method::POST && path == "/v1/messages"
29    }
30
31    fn upstream_base(&self) -> &Url {
32        &self.upstream
33    }
34
35    fn request_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
36        let mut surfaces = Vec::new();
37        if let Value::Object(root) = body {
38            for (key, value) in root {
39                match key.as_str() {
40                    "system" => push_text_blocks(&mut surfaces, "system", value),
41                    "messages" => {
42                        if let Value::Array(messages) = value {
43                            for (index, message) in messages.iter_mut().enumerate() {
44                                if let Value::Object(message) = message {
45                                    for (message_key, message_value) in message {
46                                        if message_key == "content" {
47                                            collect_content_blocks(
48                                                &mut surfaces,
49                                                &format!("messages[{index}].content"),
50                                                message_value,
51                                            );
52                                        }
53                                    }
54                                }
55                            }
56                        }
57                    }
58                    _ => {}
59                }
60            }
61        }
62        surfaces
63    }
64
65    fn response_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<PiiSurface<'a>> {
66        let mut surfaces = Vec::new();
67        if let Value::Object(root) = body {
68            if let Some(content) = root.get_mut("content") {
69                collect_content_blocks(&mut surfaces, "content", content);
70            }
71        }
72        surfaces
73    }
74
75    fn sse_event_pii_surfaces<'a>(&self, event: &'a mut SseEvent) -> Vec<PiiSurface<'a>> {
76        let mut surfaces = Vec::new();
77        if let Value::Object(root) = &mut event.data {
78            if matches!(root.get("type"), Some(Value::String(kind)) if kind == "content_block_delta")
79            {
80                if let Some(Value::Object(delta)) = root.get_mut("delta") {
81                    for (key, value) in delta {
82                        if let Value::String(text) = value {
83                            match key.as_str() {
84                                "text" => surfaces.push(PiiSurface {
85                                    field_path: "delta.text".to_string(),
86                                    text,
87                                }),
88                                "partial_json" => surfaces.push(PiiSurface {
89                                    field_path: "delta.partial_json".to_string(),
90                                    text,
91                                }),
92                                _ => {}
93                            }
94                        }
95                    }
96                }
97            }
98        }
99        surfaces
100    }
101}
102
103fn collect_content_blocks<'a>(
104    surfaces: &mut Vec<PiiSurface<'a>>,
105    prefix: &str,
106    content: &'a mut Value,
107) {
108    match content {
109        Value::String(_) => push_string(surfaces, prefix, content),
110        Value::Array(blocks) => {
111            for (index, block) in blocks.iter_mut().enumerate() {
112                let Value::Object(block) = block else {
113                    continue;
114                };
115                match block.get("type").and_then(Value::as_str) {
116                    Some("text") => {
117                        if let Some(Value::String(text)) = block.get_mut("text") {
118                            surfaces.push(PiiSurface {
119                                field_path: format!("{prefix}[{index}].text"),
120                                text,
121                            });
122                        }
123                    }
124                    Some("tool_use") => {
125                        if let Some(input) = block.get_mut("input") {
126                            walk_all_strings(surfaces, format!("{prefix}[{index}].input"), input);
127                        }
128                    }
129                    Some("tool_result") => {
130                        if let Some(result) = block.get_mut("content") {
131                            collect_content_blocks(
132                                surfaces,
133                                &format!("{prefix}[{index}].content"),
134                                result,
135                            );
136                        }
137                    }
138                    Some("image") | None | Some(_) => {}
139                }
140            }
141        }
142        _ => {}
143    }
144}