Skip to main content

gproxy_protocol/adapt/generate/resources/
claude.rs

1use super::*;
2use crate::wire::{
3    DeclaredFields,
4    claude::{content as c, generate_content as cg},
5};
6
7impl<R: ResourceAccess> GenerationResources<'_, R> {
8    pub async fn claude(
9        &self,
10        input: cg::GenerateContentRequestBody,
11    ) -> Result<cg::GenerateContentRequestBody, TransformError> {
12        let mut input = input.into_declared();
13        let mut budget = self.budget();
14        for message in &mut input.messages {
15            if let c::MessageContent::Blocks(blocks) = &mut message.content {
16                for block in blocks {
17                    match block {
18                        c::ContentBlock::Image(image) => budget.image(image).await?,
19                        c::ContentBlock::Document(doc) => budget.document(doc).await?,
20                        c::ContentBlock::ToolResult(result) => {
21                            if let Some(c::ToolResultContent::Blocks(blocks)) = &mut result.content
22                            {
23                                for block in blocks {
24                                    match block {
25                                        c::ToolResultContentBlock::Image(image) => {
26                                            budget.image(image).await?
27                                        }
28                                        c::ToolResultContentBlock::Document(doc) => {
29                                            budget.document(doc).await?
30                                        }
31                                        _ => {}
32                                    }
33                                }
34                            }
35                        }
36                        _ => {}
37                    }
38                }
39            }
40        }
41        Ok(input)
42    }
43}
44
45impl<R: ResourceAccess> Budget<'_, '_, R> {
46    async fn image(&mut self, image: &mut c::ImageBlock) -> Result<(), TransformError> {
47        let reference = match &image.source {
48            c::ImageSource::Url(source) => ResourceReference::Url(source.url.clone()),
49            c::ImageSource::File(source) => ResourceReference::Id(source.file_id.clone()),
50            c::ImageSource::Base64(_) => return Ok(()),
51        };
52        let media = self.read(reference, true).await?;
53        let mime = match media.mime.as_str() {
54            "image/png" => c::ImageMediaType::Png,
55            "image/jpeg" => c::ImageMediaType::Jpeg,
56            "image/gif" => c::ImageMediaType::Gif,
57            "image/webp" => c::ImageMediaType::Webp,
58            _ => {
59                return Err(TransformError::unsupported(
60                    "image.mime",
61                    "Claude has no native image variant for this MIME",
62                ));
63            }
64        };
65        image.source = c::ImageSource::Base64(
66            c::Base64Source::builder(STANDARD.encode(media.bytes), mime).build(),
67        );
68        Ok(())
69    }
70    async fn document(&mut self, doc: &mut c::DocumentBlock) -> Result<(), TransformError> {
71        let reference = match &mut doc.source {
72            c::DocumentSource::Url(source) => ResourceReference::Url(source.url.clone()),
73            c::DocumentSource::File(source) => ResourceReference::Id(source.file_id.clone()),
74            c::DocumentSource::Content(source) => {
75                if let c::DocumentContent::Blocks(blocks) = &mut source.content {
76                    for block in blocks {
77                        if let c::DocumentContentBlock::Image(image) = block {
78                            self.image(image).await?;
79                        }
80                    }
81                }
82                return Ok(());
83            }
84            _ => return Ok(()),
85        };
86        let media = self.read(reference, false).await?;
87        doc.source = match media.mime.as_str() {
88            "application/pdf" => {
89                if !media.bytes.starts_with(b"%PDF-") {
90                    return Err(TransformError::invalid_result(
91                        "document.mime",
92                        "PDF metadata does not match bytes",
93                    ));
94                }
95                c::DocumentSource::Base64(
96                    c::PdfBase64Source::builder(STANDARD.encode(media.bytes), c::PdfMediaType::Pdf)
97                        .build(),
98                )
99            }
100            "text/plain" => c::DocumentSource::Text(
101                c::PlainTextSource::builder(
102                    String::from_utf8(media.bytes.to_vec()).map_err(|e| {
103                        TransformError::invalid_result("document.utf8", e.to_string())
104                    })?,
105                    c::TextMediaType::Plain,
106                )
107                .build(),
108            ),
109            _ => {
110                return Err(TransformError::unsupported(
111                    "document.mime",
112                    "Claude document requires PDF or plain text",
113                ));
114            }
115        };
116        Ok(())
117    }
118}