potato_type/anthropic/v1/
request.rs

1use crate::anthropic::v1::response::{
2    ResponseContentBlockInner, TextCitation, WebSearchToolResultBlockContent,
3};
4use crate::common::get_image_media_types;
5use crate::google::v1::generate::request::{DataNum, GeminiContent, Part};
6use crate::openai::v1::chat::request::{ChatMessage, ContentPart, TextContentPart};
7use crate::prompt::builder::ProviderRequest;
8use crate::prompt::MessageNum;
9use crate::prompt::ModelSettings;
10use crate::tools::AgentToolDefinition;
11use crate::traits::get_var_regex;
12use crate::traits::{MessageConversion, MessageFactory, PromptMessageExt, RequestAdapter};
13use crate::TypeError;
14use crate::{Provider, SettingsType};
15use potato_util::{PyHelperFuncs, UtilError};
16use pyo3::prelude::*;
17use pyo3::types::PyList;
18use pyo3::IntoPyObjectExt;
19use pythonize::{depythonize, pythonize};
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22use std::collections::HashSet;
23
24/// common content types used in Anthropic messages
25pub const BASE64_TYPE: &str = "base64";
26pub const URL_TYPE: &str = "url";
27pub const EPHEMERAL_TYPE: &str = "ephemeral";
28pub const IMAGE_TYPE: &str = "image";
29pub const TEXT_TYPE: &str = "text";
30pub const DOCUMENT_TYPE: &str = "document";
31pub const DOCUMENT_BASE64_PDF_TYPE: &str = "application/pdf";
32pub const DOCUMENT_PLAIN_TEXT_TYPE: &str = "text/plain";
33pub const WEB_SEARCH_RESULT_TYPE: &str = "web_search_result";
34pub const SEARCH_TYPE: &str = "search_result";
35pub const THINKING_TYPE: &str = "thinking";
36pub const REDACTED_THINKING_TYPE: &str = "redacted_thinking";
37pub const TOOL_USE_TYPE: &str = "tool_use";
38pub const TOOL_RESULT_TYPE: &str = "tool_result";
39pub const WEB_SEARCH_TOOL_RESULT_TYPE: &str = "web_search_tool_result";
40pub const SERVER_TOOL_USE_TYPE: &str = "server_tool_use";
41
42// Citation type constants
43pub const CHAR_LOCATION_TYPE: &str = "char_location";
44pub const PAGE_LOCATION_TYPE: &str = "page_location";
45pub const CONTENT_BLOCK_LOCATION_TYPE: &str = "content_block_location";
46pub const WEB_SEARCH_RESULT_LOCATION_TYPE: &str = "web_search_result_location";
47pub const SEARCH_RESULT_LOCATION_TYPE: &str = "search_result_location";
48pub const WEB_SEARCH_TOOL_RESULT_ERROR_TYPE: &str = "web_search_tool_result_error";
49
50#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
51#[pyclass(from_py_object)]
52pub struct CitationCharLocationParam {
53    #[pyo3(get, set)]
54    pub cited_text: String,
55    #[pyo3(get, set)]
56    pub document_index: i32,
57    #[pyo3(get, set)]
58    pub document_title: String,
59    #[pyo3(get, set)]
60    pub end_char_index: i32,
61    #[pyo3(get, set)]
62    pub start_char_index: i32,
63    #[pyo3(get)]
64    #[serde(rename = "type")]
65    pub r#type: String,
66}
67
68#[pymethods]
69impl CitationCharLocationParam {
70    #[new]
71    pub fn new(
72        cited_text: String,
73        document_index: i32,
74        document_title: String,
75        end_char_index: i32,
76        start_char_index: i32,
77    ) -> Self {
78        Self {
79            cited_text,
80            document_index,
81            document_title,
82            end_char_index,
83            start_char_index,
84            r#type: CHAR_LOCATION_TYPE.to_string(),
85        }
86    }
87}
88
89#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
90#[pyclass(from_py_object)]
91pub struct CitationPageLocationParam {
92    #[pyo3(get, set)]
93    pub cited_text: String,
94    #[pyo3(get, set)]
95    pub document_index: i32,
96    #[pyo3(get, set)]
97    pub document_title: String,
98    #[pyo3(get, set)]
99    pub end_page_number: i32,
100    #[pyo3(get, set)]
101    pub start_page_number: i32,
102    #[pyo3(get)]
103    #[serde(rename = "type")]
104    pub r#type: String,
105}
106
107#[pymethods]
108impl CitationPageLocationParam {
109    #[new]
110    pub fn new(
111        cited_text: String,
112        document_index: i32,
113        document_title: String,
114        end_page_number: i32,
115        start_page_number: i32,
116    ) -> Self {
117        Self {
118            cited_text,
119            document_index,
120            document_title,
121            end_page_number,
122            start_page_number,
123            r#type: PAGE_LOCATION_TYPE.to_string(),
124        }
125    }
126}
127
128#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
129#[pyclass(from_py_object)]
130pub struct CitationContentBlockLocationParam {
131    #[pyo3(get, set)]
132    pub cited_text: String,
133    #[pyo3(get, set)]
134    pub document_index: i32,
135    #[pyo3(get, set)]
136    pub document_title: String,
137    #[pyo3(get, set)]
138    pub end_block_index: i32,
139    #[pyo3(get, set)]
140    pub start_block_index: i32,
141    #[pyo3(get)]
142    #[serde(rename = "type")]
143    pub r#type: String,
144}
145
146#[pymethods]
147impl CitationContentBlockLocationParam {
148    #[new]
149    pub fn new(
150        cited_text: String,
151        document_index: i32,
152        document_title: String,
153        end_block_index: i32,
154        start_block_index: i32,
155    ) -> Self {
156        Self {
157            cited_text,
158            document_index,
159            document_title,
160            end_block_index,
161            start_block_index,
162            r#type: CONTENT_BLOCK_LOCATION_TYPE.to_string(),
163        }
164    }
165}
166
167#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
168#[pyclass(from_py_object)]
169pub struct CitationWebSearchResultLocationParam {
170    #[pyo3(get, set)]
171    pub cited_text: String,
172    #[pyo3(get, set)]
173    pub encrypted_index: String,
174    #[pyo3(get, set)]
175    pub title: String,
176    #[pyo3(get)]
177    #[serde(rename = "type")]
178    pub r#type: String,
179    #[pyo3(get, set)]
180    pub url: String,
181}
182
183#[pymethods]
184impl CitationWebSearchResultLocationParam {
185    #[new]
186    pub fn new(cited_text: String, encrypted_index: String, title: String, url: String) -> Self {
187        Self {
188            cited_text,
189            encrypted_index,
190            title,
191            r#type: WEB_SEARCH_RESULT_LOCATION_TYPE.to_string(),
192            url,
193        }
194    }
195}
196
197#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
198#[pyclass(from_py_object)]
199pub struct CitationSearchResultLocationParam {
200    #[pyo3(get, set)]
201    pub cited_text: String,
202    #[pyo3(get, set)]
203    pub end_block_index: i32,
204    #[pyo3(get, set)]
205    pub search_result_index: i32,
206    #[pyo3(get, set)]
207    pub source: String,
208    #[pyo3(get, set)]
209    pub start_block_index: i32,
210    #[pyo3(get, set)]
211    pub title: String,
212    #[pyo3(get)]
213    #[serde(rename = "type")]
214    pub r#type: String,
215}
216
217#[pymethods]
218impl CitationSearchResultLocationParam {
219    #[new]
220    pub fn new(
221        cited_text: String,
222        end_block_index: i32,
223        search_result_index: i32,
224        source: String,
225        start_block_index: i32,
226        title: String,
227    ) -> Self {
228        Self {
229            cited_text,
230            end_block_index,
231            search_result_index,
232            source,
233            start_block_index,
234            title,
235            r#type: SEARCH_RESULT_LOCATION_TYPE.to_string(),
236        }
237    }
238}
239
240/// Untagged enum for internal Rust usage - serializes without wrapper
241#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
242#[serde(untagged)]
243pub enum TextCitationParam {
244    CharLocation(CitationCharLocationParam),
245    PageLocation(CitationPageLocationParam),
246    ContentBlockLocation(CitationContentBlockLocationParam),
247    WebSearchResultLocation(CitationWebSearchResultLocationParam),
248    SearchResultLocation(CitationSearchResultLocationParam),
249}
250
251#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
252#[pyclass(from_py_object)]
253pub struct TextBlockParam {
254    #[pyo3(get, set)]
255    pub text: String,
256    #[serde(skip_serializing_if = "Option::is_none")]
257    #[pyo3(get, set)]
258    pub cache_control: Option<CacheControl>,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub citations: Option<TextCitationParam>,
261    #[pyo3(get)]
262    #[serde(rename = "type")]
263    pub r#type: String,
264}
265
266fn parse_text_citation(cit: &Bound<'_, PyAny>) -> Result<TextCitationParam, TypeError> {
267    if cit.is_instance_of::<CitationCharLocationParam>() {
268        Ok(TextCitationParam::CharLocation(
269            cit.extract::<CitationCharLocationParam>()?,
270        ))
271    } else if cit.is_instance_of::<CitationPageLocationParam>() {
272        Ok(TextCitationParam::PageLocation(
273            cit.extract::<CitationPageLocationParam>()?,
274        ))
275    } else if cit.is_instance_of::<CitationContentBlockLocationParam>() {
276        Ok(TextCitationParam::ContentBlockLocation(
277            cit.extract::<CitationContentBlockLocationParam>()?,
278        ))
279    } else if cit.is_instance_of::<CitationWebSearchResultLocationParam>() {
280        Ok(TextCitationParam::WebSearchResultLocation(
281            cit.extract::<CitationWebSearchResultLocationParam>()?,
282        ))
283    } else if cit.is_instance_of::<CitationSearchResultLocationParam>() {
284        Ok(TextCitationParam::SearchResultLocation(
285            cit.extract::<CitationSearchResultLocationParam>()?,
286        ))
287    } else {
288        Err(TypeError::InvalidInput(
289            "Invalid citation type provided".to_string(),
290        ))
291    }
292}
293#[pymethods]
294impl TextBlockParam {
295    #[new]
296    #[pyo3(signature = (text, cache_control=None, citations=None))]
297    pub fn new(
298        text: String,
299        cache_control: Option<CacheControl>,
300        citations: Option<&Bound<'_, PyAny>>,
301    ) -> Result<Self, TypeError> {
302        let citations = if let Some(cit) = citations {
303            Some(parse_text_citation(cit)?)
304        } else {
305            None
306        };
307        Ok(Self {
308            text,
309            cache_control,
310            citations,
311            r#type: TEXT_TYPE.to_string(),
312        })
313    }
314}
315
316impl TextBlockParam {
317    pub fn new_rs(
318        text: String,
319        cache_control: Option<CacheControl>,
320        citations: Option<TextCitationParam>,
321    ) -> Self {
322        Self {
323            text,
324            cache_control,
325            citations,
326            r#type: TEXT_TYPE.to_string(),
327        }
328    }
329}
330
331#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
332#[pyclass(from_py_object)]
333pub struct Base64ImageSource {
334    #[pyo3(get, set)]
335    pub media_type: String,
336    #[pyo3(get, set)]
337    pub data: String,
338    #[pyo3(get)]
339    #[serde(rename = "type")]
340    pub r#type: String,
341}
342
343#[pymethods]
344impl Base64ImageSource {
345    #[new]
346    pub fn new(media_type: String, data: String) -> Result<Self, TypeError> {
347        // confirm media_type is an image type, otherwise raise error
348        if !get_image_media_types().contains(media_type.as_str()) {
349            return Err(TypeError::InvalidMediaType(media_type));
350        }
351        Ok(Self {
352            media_type,
353            data,
354            r#type: BASE64_TYPE.to_string(),
355        })
356    }
357}
358
359#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
360#[pyclass(from_py_object)]
361pub struct UrlImageSource {
362    #[pyo3(get, set)]
363    pub url: String,
364    #[pyo3(get)]
365    #[serde(rename = "type")]
366    pub r#type: String,
367}
368
369#[pymethods]
370impl UrlImageSource {
371    #[new]
372    pub fn new(url: String) -> Self {
373        Self {
374            url,
375            r#type: URL_TYPE.to_string(),
376        }
377    }
378}
379
380#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
381#[serde(untagged)] // we need to strip serde type ref
382pub enum ImageSource {
383    Base64(Base64ImageSource),
384    Url(UrlImageSource),
385}
386
387#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
388#[pyclass(from_py_object)]
389pub struct ImageBlockParam {
390    pub source: ImageSource,
391    #[serde(skip_serializing_if = "Option::is_none")]
392    #[pyo3(get, set)]
393    pub cache_control: Option<CacheControl>,
394    #[pyo3(get)]
395    #[serde(rename = "type")]
396    pub r#type: String,
397}
398
399#[pymethods]
400impl ImageBlockParam {
401    #[new]
402    #[pyo3(signature = (source, cache_control=None))]
403    pub fn new(
404        source: &Bound<'_, PyAny>,
405        cache_control: Option<CacheControl>,
406    ) -> Result<Self, TypeError> {
407        let source: ImageSource = if source.is_instance_of::<Base64ImageSource>() {
408            ImageSource::Base64(source.extract::<Base64ImageSource>()?)
409        } else {
410            ImageSource::Url(source.extract::<UrlImageSource>()?)
411        };
412        Ok(Self {
413            source,
414            cache_control,
415            r#type: IMAGE_TYPE.to_string(),
416        })
417    }
418
419    #[getter]
420    pub fn source<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
421        match &self.source {
422            ImageSource::Base64(base64) => {
423                let py_obj = base64.clone().into_bound_py_any(py)?;
424                Ok(py_obj.clone())
425            }
426            ImageSource::Url(url) => {
427                let py_obj = url.clone().into_bound_py_any(py)?;
428                Ok(py_obj.clone())
429            }
430        }
431    }
432}
433
434#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
435#[pyclass(from_py_object)]
436pub struct Base64PDFSource {
437    #[pyo3(get, set)]
438    pub media_type: String,
439    #[pyo3(get, set)]
440    pub data: String,
441    #[pyo3(get)]
442    #[serde(rename = "type")]
443    pub r#type: String,
444}
445
446#[pymethods]
447impl Base64PDFSource {
448    #[new]
449    pub fn new(data: String) -> Result<Self, TypeError> {
450        Ok(Self {
451            media_type: DOCUMENT_BASE64_PDF_TYPE.to_string(),
452            data,
453            r#type: BASE64_TYPE.to_string(),
454        })
455    }
456}
457
458#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
459#[pyclass(from_py_object)]
460pub struct UrlPDFSource {
461    #[pyo3(get, set)]
462    pub url: String,
463    #[pyo3(get)]
464    #[serde(rename = "type")]
465    pub r#type: String,
466}
467
468#[pymethods]
469impl UrlPDFSource {
470    #[new]
471    pub fn new(url: String) -> Self {
472        Self {
473            url,
474            r#type: URL_TYPE.to_string(),
475        }
476    }
477}
478
479#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
480#[pyclass(from_py_object)]
481pub struct PlainTextSource {
482    #[pyo3(get, set)]
483    pub media_type: String,
484    #[pyo3(get, set)]
485    pub data: String,
486    #[pyo3(get)]
487    #[serde(rename = "type")]
488    pub r#type: String,
489}
490
491#[pymethods]
492impl PlainTextSource {
493    #[new]
494    pub fn new(data: String) -> Self {
495        Self {
496            media_type: DOCUMENT_PLAIN_TEXT_TYPE.to_string(),
497            data,
498            r#type: TEXT_TYPE.to_string(),
499        }
500    }
501}
502
503#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
504#[pyclass(from_py_object)]
505pub struct CitationsConfigParams {
506    #[pyo3(get, set)]
507    pub enabled: Option<bool>,
508}
509
510#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
511#[serde(untagged)]
512pub enum DocumentSource {
513    Base64(Base64PDFSource),
514    Url(UrlPDFSource),
515    Text(PlainTextSource),
516}
517
518#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
519#[pyclass(from_py_object)]
520pub struct DocumentBlockParam {
521    pub source: DocumentSource,
522    #[serde(skip_serializing_if = "Option::is_none")]
523    #[pyo3(get, set)]
524    pub cache_control: Option<CacheControl>,
525    #[serde(skip_serializing_if = "Option::is_none")]
526    #[pyo3(get, set)]
527    pub title: Option<String>,
528    #[serde(skip_serializing_if = "Option::is_none")]
529    #[pyo3(get, set)]
530    pub context: Option<String>,
531    #[serde(rename = "type")]
532    #[pyo3(get, set)]
533    pub r#type: String,
534    #[serde(skip_serializing_if = "Option::is_none")]
535    #[pyo3(get, set)]
536    pub citations: Option<CitationsConfigParams>,
537}
538
539#[pymethods]
540impl DocumentBlockParam {
541    #[new]
542    #[pyo3(signature = (source, cache_control=None, title=None, context=None, citations=None))]
543    pub fn new(
544        source: &Bound<'_, PyAny>,
545        cache_control: Option<CacheControl>,
546        title: Option<String>,
547        context: Option<String>,
548        citations: Option<CitationsConfigParams>,
549    ) -> Result<Self, TypeError> {
550        let source: DocumentSource = if source.is_instance_of::<Base64PDFSource>() {
551            DocumentSource::Base64(source.extract::<Base64PDFSource>()?)
552        } else if source.is_instance_of::<UrlPDFSource>() {
553            DocumentSource::Url(source.extract::<UrlPDFSource>()?)
554        } else {
555            DocumentSource::Text(source.extract::<PlainTextSource>()?)
556        };
557
558        Ok(Self {
559            source,
560            cache_control,
561            title,
562            context,
563            r#type: DOCUMENT_TYPE.to_string(),
564            citations,
565        })
566    }
567}
568
569#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
570#[pyclass(from_py_object)]
571pub struct SearchResultBlockParam {
572    #[pyo3(get, set)]
573    pub content: Vec<TextBlockParam>,
574    #[pyo3(get, set)]
575    pub source: String,
576    #[pyo3(get, set)]
577    pub title: String,
578    #[serde(rename = "type")]
579    #[pyo3(get, set)]
580    pub r#type: String,
581    #[pyo3(get, set)]
582    pub cache_control: Option<CacheControl>,
583    #[pyo3(get, set)]
584    pub citations: Option<CitationsConfigParams>,
585}
586
587#[pymethods]
588impl SearchResultBlockParam {
589    #[new]
590    #[pyo3(signature = (content, source, title, cache_control=None, citations=None))]
591    pub fn new(
592        content: Vec<TextBlockParam>,
593        source: String,
594        title: String,
595        cache_control: Option<CacheControl>,
596        citations: Option<CitationsConfigParams>,
597    ) -> Self {
598        Self {
599            content,
600            source,
601            title,
602            r#type: SEARCH_TYPE.to_string(),
603            cache_control,
604            citations,
605        }
606    }
607}
608
609#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
610#[pyclass(from_py_object)]
611pub struct ThinkingBlockParam {
612    #[pyo3(get, set)]
613    pub thinking: String,
614    #[pyo3(get, set)]
615    pub signature: Option<String>,
616    #[pyo3(get, set)]
617    #[serde(rename = "type")]
618    pub r#type: String,
619}
620
621#[pymethods]
622impl ThinkingBlockParam {
623    #[new]
624    #[pyo3(signature = (thinking, signature=None))]
625    pub fn new(thinking: String, signature: Option<String>) -> Self {
626        Self {
627            thinking,
628            signature,
629            r#type: THINKING_TYPE.to_string(),
630        }
631    }
632}
633
634#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
635#[pyclass(from_py_object)]
636pub struct RedactedThinkingBlockParam {
637    #[pyo3(get, set)]
638    pub data: String,
639    #[pyo3(get, set)]
640    #[serde(rename = "type")]
641    pub r#type: String,
642}
643
644#[pymethods]
645impl RedactedThinkingBlockParam {
646    #[new]
647    pub fn new(data: String) -> Self {
648        Self {
649            data,
650            r#type: REDACTED_THINKING_TYPE.to_string(),
651        }
652    }
653}
654
655#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
656#[pyclass(from_py_object)]
657pub struct ToolUseBlockParam {
658    #[pyo3(get, set)]
659    pub id: String,
660    #[pyo3(get, set)]
661    pub name: String,
662    pub input: Value,
663    #[serde(skip_serializing_if = "Option::is_none")]
664    #[pyo3(get, set)]
665    pub cache_control: Option<CacheControl>,
666    #[pyo3(get, set)]
667    #[serde(rename = "type")]
668    pub r#type: String,
669}
670
671#[pymethods]
672impl ToolUseBlockParam {
673    #[new]
674    #[pyo3(signature = (id, name, input, cache_control=None))]
675    pub fn new(
676        id: String,
677        name: String,
678        input: &Bound<'_, PyAny>,
679        cache_control: Option<CacheControl>,
680    ) -> Result<Self, TypeError> {
681        let input_value = depythonize(input)?;
682        Ok(Self {
683            id,
684            name,
685            input: input_value,
686            cache_control,
687            r#type: TOOL_USE_TYPE.to_string(),
688        })
689    }
690
691    #[getter]
692    pub fn input<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
693        let py_dict = pythonize(py, &self.input)?;
694        Ok(py_dict)
695    }
696}
697
698#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
699#[serde(untagged)]
700pub enum ToolResultContentEnum {
701    Text(Vec<TextBlockParam>),
702    Image(Vec<ImageBlockParam>),
703    Document(Vec<DocumentBlockParam>),
704    SearchResult(Vec<SearchResultBlockParam>),
705}
706
707#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
708#[pyclass(from_py_object)]
709pub struct ToolResultBlockParam {
710    #[pyo3(get, set)]
711    pub tool_use_id: String,
712    #[serde(skip_serializing_if = "Option::is_none")]
713    pub is_error: Option<bool>,
714    #[serde(skip_serializing_if = "Option::is_none")]
715    #[pyo3(get, set)]
716    pub cache_control: Option<CacheControl>,
717    #[pyo3(get, set)]
718    #[serde(rename = "type")]
719    pub r#type: String,
720    #[serde(skip_serializing_if = "Option::is_none")]
721    pub content: Option<ToolResultContentEnum>,
722}
723
724/// Helper function to extract all blocks of a specific type
725fn extract_all_blocks<T>(blocks: Vec<Bound<'_, PyAny>>) -> Result<Vec<T>, TypeError>
726where
727    T: for<'a, 'py> FromPyObject<'a, 'py>,
728{
729    blocks
730        .into_iter()
731        .map(|block| {
732            block
733                .extract::<T>()
734                .map_err(|_| TypeError::Error("Failed to extract block".to_string()))
735        })
736        .collect()
737}
738
739#[pymethods]
740impl ToolResultBlockParam {
741    #[new]
742    #[pyo3(signature = (
743        tool_use_id,
744        is_error=None,
745        cache_control=None,
746        content=None
747    ))]
748    pub fn new(
749        tool_use_id: String,
750        is_error: Option<bool>,
751        cache_control: Option<CacheControl>,
752        content: Option<Vec<Bound<'_, PyAny>>>,
753    ) -> Result<Self, TypeError> {
754        let content_enum = match content {
755            None => None,
756            Some(blocks) if blocks.is_empty() => None,
757
758            Some(blocks) => {
759                let first_block = &blocks[0];
760
761                if first_block.is_instance_of::<TextBlockParam>() {
762                    Some(ToolResultContentEnum::Text(extract_all_blocks(blocks)?))
763                } else if first_block.is_instance_of::<ImageBlockParam>() {
764                    Some(ToolResultContentEnum::Image(extract_all_blocks(blocks)?))
765                } else if first_block.is_instance_of::<DocumentBlockParam>() {
766                    Some(ToolResultContentEnum::Document(extract_all_blocks(blocks)?))
767                } else if first_block.is_instance_of::<SearchResultBlockParam>() {
768                    Some(ToolResultContentEnum::SearchResult(extract_all_blocks(
769                        blocks,
770                    )?))
771                } else {
772                    return Err(TypeError::InvalidInput(
773                        "Unsupported content block type".to_string(),
774                    ));
775                }
776            }
777        };
778
779        Ok(Self {
780            tool_use_id,
781            is_error,
782            cache_control,
783            r#type: TOOL_RESULT_TYPE.to_string(),
784            content: content_enum,
785        })
786    }
787
788    #[getter]
789    pub fn content<'py>(&self, py: Python<'py>) -> Result<Option<Bound<'py, PyAny>>, TypeError> {
790        match &self.content {
791            None => Ok(None),
792            Some(ToolResultContentEnum::Text(blocks)) => {
793                let py_list = blocks
794                    .iter()
795                    .map(|block| block.clone().into_bound_py_any(py))
796                    .collect::<Result<Vec<_>, _>>()?;
797                Ok(Some(py_list.into_bound_py_any(py)?))
798            }
799            Some(ToolResultContentEnum::Image(blocks)) => {
800                let py_list = blocks
801                    .iter()
802                    .map(|block| block.clone().into_bound_py_any(py))
803                    .collect::<Result<Vec<_>, _>>()?;
804                Ok(Some(py_list.into_bound_py_any(py)?))
805            }
806            Some(ToolResultContentEnum::Document(blocks)) => {
807                let py_list = blocks
808                    .iter()
809                    .map(|block| block.clone().into_bound_py_any(py))
810                    .collect::<Result<Vec<_>, _>>()?;
811                Ok(Some(py_list.into_bound_py_any(py)?))
812            }
813            Some(ToolResultContentEnum::SearchResult(blocks)) => {
814                let py_list = blocks
815                    .iter()
816                    .map(|block| block.clone().into_bound_py_any(py))
817                    .collect::<Result<Vec<_>, _>>()?;
818                Ok(Some(py_list.into_bound_py_any(py)?))
819            }
820        }
821    }
822}
823
824#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
825#[pyclass(from_py_object)]
826pub struct ServerToolUseBlockParam {
827    #[pyo3(get, set)]
828    pub id: String,
829    #[pyo3(get, set)]
830    pub name: String,
831    pub input: Value,
832    #[serde(skip_serializing_if = "Option::is_none")]
833    pub cache_control: Option<CacheControl>,
834    #[pyo3(get, set)]
835    #[serde(rename = "type")]
836    pub r#type: String,
837}
838
839#[pymethods]
840impl ServerToolUseBlockParam {
841    #[new]
842    #[pyo3(signature = (id, name, input, cache_control=None))]
843    pub fn new(
844        id: String,
845        name: String,
846        input: &Bound<'_, PyAny>,
847        cache_control: Option<CacheControl>,
848    ) -> Result<Self, TypeError> {
849        let input_value = depythonize(input)?;
850        Ok(Self {
851            id,
852            name,
853            input: input_value,
854            cache_control,
855            r#type: SERVER_TOOL_USE_TYPE.to_string(),
856        })
857    }
858    #[getter]
859    pub fn input<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
860        let py_dict = pythonize(py, &self.input)?;
861        Ok(py_dict)
862    }
863}
864
865#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
866#[pyclass(from_py_object)]
867pub struct WebSearchResultBlockParam {
868    #[pyo3(get, set)]
869    pub encrypted_content: String,
870    #[pyo3(get, set)]
871    pub title: String,
872    #[pyo3(get, set)]
873    pub url: String,
874    #[pyo3(get, set)]
875    pub page_agent: Option<String>,
876    #[pyo3(get, set)]
877    #[serde(rename = "type")]
878    pub r#type: String,
879}
880
881#[pymethods]
882impl WebSearchResultBlockParam {
883    #[new]
884    #[pyo3(signature = (encrypted_content, title, url, page_agent=None))]
885    pub fn new(
886        encrypted_content: String,
887        title: String,
888        url: String,
889        page_agent: Option<String>,
890    ) -> Self {
891        Self {
892            encrypted_content,
893            title,
894            url,
895            page_agent,
896            r#type: WEB_SEARCH_RESULT_TYPE.to_string(),
897        }
898    }
899}
900
901#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
902#[pyclass(from_py_object)]
903pub struct WebSearchToolResultBlockParam {
904    #[pyo3(get, set)]
905    pub tool_use_id: String,
906    #[pyo3(get, set)]
907    pub content: Vec<WebSearchResultBlockParam>,
908    #[serde(skip_serializing_if = "Option::is_none")]
909    #[pyo3(get, set)]
910    pub cache_control: Option<CacheControl>,
911    #[pyo3(get, set)]
912    #[serde(rename = "type")]
913    pub r#type: String,
914}
915
916#[pymethods]
917impl WebSearchToolResultBlockParam {
918    #[new]
919    #[pyo3(signature = (tool_use_id, content, cache_control=None))]
920    pub fn new(
921        tool_use_id: String,
922        content: Vec<WebSearchResultBlockParam>,
923        cache_control: Option<CacheControl>,
924    ) -> Self {
925        Self {
926            tool_use_id,
927            content,
928            cache_control,
929            r#type: WEB_SEARCH_TOOL_RESULT_TYPE.to_string(),
930        }
931    }
932}
933
934#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
935#[serde(untagged)]
936pub(crate) enum ContentBlock {
937    Text(TextBlockParam),
938    Image(ImageBlockParam),
939    Document(DocumentBlockParam),
940    SearchResult(SearchResultBlockParam),
941    Thinking(ThinkingBlockParam),
942    RedactedThinking(RedactedThinkingBlockParam),
943    ToolUse(ToolUseBlockParam),
944    ToolResult(ToolResultBlockParam),
945    ServerToolUse(ServerToolUseBlockParam),
946    WebSearchResult(WebSearchResultBlockParam),
947}
948
949macro_rules! try_extract_content_block {
950    ($block:expr, $($variant:ident => $type:ty),+ $(,)?) => {{
951        $(
952            if $block.is_instance_of::<$type>() {
953                return Ok(Self {
954                    inner: ContentBlock::$variant($block.extract::<$type>()?),
955                });
956            }
957        )+
958        return Err(TypeError::InvalidInput(
959            "Unsupported content block type".to_string(),
960        ));
961    }};
962}
963
964#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
965pub struct ContentBlockParam {
966    #[serde(flatten)]
967    pub(crate) inner: ContentBlock,
968}
969
970impl ContentBlockParam {
971    pub fn new(block: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
972        try_extract_content_block!(
973            block,
974            Text => TextBlockParam,
975            Image => ImageBlockParam,
976            Document => DocumentBlockParam,
977            SearchResult => SearchResultBlockParam,
978            Thinking => ThinkingBlockParam,
979            RedactedThinking => RedactedThinkingBlockParam,
980            ToolUse => ToolUseBlockParam,
981            ToolResult => ToolResultBlockParam,
982            ServerToolUse => ServerToolUseBlockParam,
983            WebSearchResult => WebSearchResultBlockParam,
984        )
985    }
986
987    /// Convert the ContentBlockParam back to a PyObject
988    /// This is an acceptable clone, as this will really only be used in development/testing scenarios
989    pub fn to_pyobject<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
990        match &self.inner {
991            ContentBlock::Text(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
992            ContentBlock::Image(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
993            ContentBlock::Document(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
994            ContentBlock::SearchResult(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
995            ContentBlock::Thinking(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
996            ContentBlock::RedactedThinking(block) => {
997                Ok(block.clone().into_bound_py_any(py)?.clone())
998            }
999            ContentBlock::ToolUse(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
1000            ContentBlock::ToolResult(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
1001            ContentBlock::ServerToolUse(block) => Ok(block.clone().into_bound_py_any(py)?.clone()),
1002            ContentBlock::WebSearchResult(block) => {
1003                Ok(block.clone().into_bound_py_any(py)?.clone())
1004            }
1005        }
1006    }
1007}
1008
1009impl ContentBlockParam {
1010    /// Helper function to create a ContentBlockParam from a ResponseContentBlockInner
1011    ///
1012    /// Converts response content blocks (from API responses) into request content blocks
1013    /// that can be used in subsequent requests. This is useful for multi-turn conversations
1014    /// where the assistant's response needs to be included in the next request.
1015    ///
1016    /// # Arguments
1017    /// * `block` - A reference to a ResponseContentBlockInner from an API response
1018    ///
1019    /// # Returns
1020    /// * `Result<Self, TypeError>` - The converted ContentBlockParam or an error
1021    pub(crate) fn from_response_content_block(
1022        block: &ResponseContentBlockInner,
1023    ) -> Result<Self, TypeError> {
1024        match block {
1025            ResponseContentBlockInner::Text(text_block) => {
1026                // Convert Vec<TextCitation> to Option<TextCitationParam>
1027                // Take the first citation if available, as request format uses singular citation
1028                let citations = text_block.citations.as_ref().and_then(|cits| {
1029                    cits.first().map(|cit| match cit {
1030                        TextCitation::CharLocation(c) => {
1031                            TextCitationParam::CharLocation(CitationCharLocationParam {
1032                                cited_text: c.cited_text.clone(),
1033                                document_index: c.document_index,
1034                                document_title: c.document_title.clone(),
1035                                end_char_index: c.end_char_index,
1036                                start_char_index: c.start_char_index,
1037                                r#type: c.r#type.clone(),
1038                            })
1039                        }
1040                        TextCitation::PageLocation(c) => {
1041                            TextCitationParam::PageLocation(CitationPageLocationParam {
1042                                cited_text: c.cited_text.clone(),
1043                                document_index: c.document_index,
1044                                document_title: c.document_title.clone(),
1045                                end_page_number: c.end_page_number,
1046                                start_page_number: c.start_page_number,
1047                                r#type: c.r#type.clone(),
1048                            })
1049                        }
1050                        TextCitation::ContentBlockLocation(c) => {
1051                            TextCitationParam::ContentBlockLocation(
1052                                CitationContentBlockLocationParam {
1053                                    cited_text: c.cited_text.clone(),
1054                                    document_index: c.document_index,
1055                                    document_title: c.document_title.clone(),
1056                                    end_block_index: c.end_block_index,
1057                                    start_block_index: c.start_block_index,
1058                                    r#type: c.r#type.clone(),
1059                                },
1060                            )
1061                        }
1062                        TextCitation::WebSearchResultLocation(c) => {
1063                            TextCitationParam::WebSearchResultLocation(
1064                                CitationWebSearchResultLocationParam {
1065                                    cited_text: c.cited_text.clone(),
1066                                    encrypted_index: c.encrypted_index.clone(),
1067                                    title: c.title.clone(),
1068                                    r#type: c.r#type.clone(),
1069                                    url: c.url.clone(),
1070                                },
1071                            )
1072                        }
1073                        TextCitation::SearchResultLocation(c) => {
1074                            TextCitationParam::SearchResultLocation(
1075                                CitationSearchResultLocationParam {
1076                                    cited_text: c.cited_text.clone(),
1077                                    end_block_index: c.end_block_index,
1078                                    search_result_index: c.search_result_index,
1079                                    source: c.source.clone(),
1080                                    start_block_index: c.start_block_index,
1081                                    title: c.title.clone(),
1082                                    r#type: c.r#type.clone(),
1083                                },
1084                            )
1085                        }
1086                    })
1087                });
1088
1089                Ok(Self {
1090                    inner: ContentBlock::Text(TextBlockParam {
1091                        text: text_block.text.clone(),
1092                        cache_control: None,
1093                        citations,
1094                        r#type: text_block.r#type.clone(),
1095                    }),
1096                })
1097            }
1098            ResponseContentBlockInner::Thinking(thinking_block) => Ok(Self {
1099                inner: ContentBlock::Thinking(ThinkingBlockParam {
1100                    thinking: thinking_block.thinking.clone(),
1101                    signature: thinking_block.signature.clone(),
1102                    r#type: thinking_block.r#type.clone(),
1103                }),
1104            }),
1105            ResponseContentBlockInner::RedactedThinking(redacted_thinking_block) => Ok(Self {
1106                inner: ContentBlock::RedactedThinking(RedactedThinkingBlockParam {
1107                    data: redacted_thinking_block.data.clone(),
1108                    r#type: redacted_thinking_block.r#type.clone(),
1109                }),
1110            }),
1111            ResponseContentBlockInner::ToolUse(tool_use_block) => Ok(Self {
1112                inner: ContentBlock::ToolUse(ToolUseBlockParam {
1113                    id: tool_use_block.id.clone(),
1114                    name: tool_use_block.name.clone(),
1115                    input: tool_use_block.input.clone(),
1116                    cache_control: None,
1117                    r#type: tool_use_block.r#type.clone(),
1118                }),
1119            }),
1120            ResponseContentBlockInner::ServerToolUse(server_tool_use_block) => Ok(Self {
1121                inner: ContentBlock::ServerToolUse(ServerToolUseBlockParam {
1122                    id: server_tool_use_block.id.clone(),
1123                    name: server_tool_use_block.name.clone(),
1124                    input: server_tool_use_block.input.clone(),
1125                    cache_control: None,
1126                    r#type: server_tool_use_block.r#type.clone(),
1127                }),
1128            }),
1129            ResponseContentBlockInner::WebSearchToolResult(web_search_tool_result_block) => {
1130                match &web_search_tool_result_block.content {
1131                    WebSearchToolResultBlockContent::Results(results) => {
1132                        // Take the first result and convert it to WebSearchResultBlockParam
1133                        let first_result = results.first().ok_or_else(|| {
1134                            TypeError::InvalidInput(
1135                                "WebSearchToolResult must contain at least one result".to_string(),
1136                            )
1137                        })?;
1138
1139                        Ok(Self {
1140                            inner: ContentBlock::WebSearchResult(WebSearchResultBlockParam {
1141                                encrypted_content: first_result.encrypted_content.clone(),
1142                                title: first_result.title.clone(),
1143                                url: first_result.url.clone(),
1144                                page_agent: first_result.page_age.clone(), // Note: page_age -> page_agent mapping
1145                                r#type: first_result.r#type.clone(),
1146                            }),
1147                        })
1148                    }
1149                    WebSearchToolResultBlockContent::Error(_) => Err(TypeError::InvalidInput(
1150                        "Cannot convert WebSearchToolResult error to ContentBlockParam".to_string(),
1151                    )),
1152                }
1153            }
1154        }
1155    }
1156}
1157
1158#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1159#[pyclass(from_py_object)]
1160pub struct MessageParam {
1161    pub content: Vec<ContentBlockParam>,
1162    #[pyo3(get)]
1163    pub role: String,
1164}
1165
1166#[pymethods]
1167impl MessageParam {
1168    /// Create a new MessageParam
1169    /// Will accept either a string, Single type of ContentBlockParam, or a list of varying ContentBlockParams
1170    /// Initial logic determines initial type and processes accordingly
1171    /// If string, wraps in TextBlockParam --> ContentBlockParam
1172    /// If single param, pass to ContentBlockParam directly
1173    /// If list, process each item in list to ContentBlockParam
1174    #[new]
1175    pub fn new(content: &Bound<'_, PyAny>, role: String) -> Result<Self, TypeError> {
1176        let content: Vec<ContentBlockParam> = if content.is_instance_of::<pyo3::types::PyString>() {
1177            let text = content.extract::<String>()?;
1178            let text_block = TextBlockParam::new(text, None, None)?;
1179            let content_block =
1180                ContentBlockParam::new(&text_block.into_bound_py_any(content.py())?)?;
1181            vec![content_block]
1182        } else if content.is_instance_of::<pyo3::types::PyList>() {
1183            let content_list = content.extract::<Vec<Bound<'_, PyAny>>>()?;
1184            let mut blocks = Vec::new();
1185            for item in content_list {
1186                let content_block = ContentBlockParam::new(&item)?;
1187                blocks.push(content_block);
1188            }
1189            blocks
1190        } else {
1191            // pass single ContentBlockParam
1192            let content_block = ContentBlockParam::new(content)?;
1193            vec![content_block]
1194        };
1195
1196        Ok(Self { content, role })
1197    }
1198
1199    #[getter]
1200    fn content<'py>(&self, py: Python<'py>) -> Result<Vec<Bound<'py, PyAny>>, TypeError> {
1201        self.content
1202            .iter()
1203            .map(|block| block.to_pyobject(py))
1204            .collect()
1205    }
1206
1207    #[pyo3(name = "bind")]
1208    fn bind_py(&self, name: &str, value: &str) -> Result<Self, TypeError> {
1209        self.bind(name, value)
1210    }
1211
1212    #[pyo3(name = "bind_mut")]
1213    fn bind_mut_py(&mut self, name: &str, value: &str) -> Result<(), TypeError> {
1214        self.bind_mut(name, value)
1215    }
1216
1217    pub fn __str__(&self) -> String {
1218        PyHelperFuncs::__str__(self)
1219    }
1220
1221    pub fn model_dump<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
1222        // iterate over each field in model_settings and add to the dict if it is not None
1223        let json = serde_json::to_value(self)?;
1224        Ok(pythonize(py, &json)?)
1225    }
1226
1227    // Return the text content from the first content part that is text
1228    #[getter]
1229    pub fn text(&self) -> String {
1230        self.content
1231            .iter()
1232            .find_map(|part| {
1233                if let ContentBlock::Text(text_part) = &part.inner {
1234                    Some(text_part.text.clone())
1235                } else {
1236                    None
1237                }
1238            })
1239            .unwrap_or_default()
1240    }
1241}
1242
1243impl PromptMessageExt for MessageParam {
1244    fn bind_mut(&mut self, name: &str, value: &str) -> Result<(), TypeError> {
1245        let regex = get_var_regex();
1246
1247        for part in &mut self.content {
1248            if let ContentBlock::Text(text_part) = &mut part.inner {
1249                text_part.text = regex
1250                    .replace_all(&text_part.text, |caps: &regex::Captures| {
1251                        if &caps[1] == name {
1252                            value.to_string()
1253                        } else {
1254                            caps.get(0)
1255                                .map_or_else(|| "".to_string(), |m| m.as_str().to_string())
1256                        }
1257                    })
1258                    .to_string();
1259            }
1260        }
1261
1262        Ok(())
1263    }
1264
1265    fn bind(&self, name: &str, value: &str) -> Result<Self, TypeError>
1266    where
1267        Self: Sized,
1268    {
1269        let mut new_message = self.clone();
1270        new_message.bind_mut(name, value)?;
1271        Ok(new_message)
1272    }
1273
1274    fn extract_variables(&self) -> Vec<String> {
1275        let mut variables = HashSet::new();
1276
1277        // Lazily initialize regex to avoid recompilation
1278        let regex = get_var_regex();
1279
1280        // Extract variables from all text content parts
1281        for part in &self.content {
1282            if let ContentBlock::Text(text_part) = &part.inner {
1283                for captures in regex.captures_iter(&text_part.text) {
1284                    if let Some(name) = captures.get(1) {
1285                        variables.insert(name.as_str().to_string());
1286                    }
1287                }
1288            }
1289        }
1290
1291        // Convert HashSet to Vec for return
1292        variables.into_iter().collect()
1293    }
1294
1295    fn from_text(content: String, role: &str) -> Result<Self, TypeError> {
1296        Ok(Self {
1297            role: role.to_string(),
1298            content: vec![ContentBlockParam {
1299                inner: ContentBlock::Text(TextBlockParam::new_rs(content, None, None)),
1300            }],
1301        })
1302    }
1303
1304    fn extract_media_variables(&self) -> Vec<String> {
1305        let mut out = HashSet::new();
1306        let regex = crate::traits::get_media_regex();
1307        for part in &self.content {
1308            if let ContentBlock::Text(t) = &part.inner {
1309                for cap in regex.captures_iter(&t.text) {
1310                    if let Some(name) = cap.get(1) {
1311                        out.insert(name.as_str().to_string());
1312                    }
1313                }
1314            }
1315        }
1316        out.into_iter().collect()
1317    }
1318
1319    fn split_media_placeholders(&mut self) -> Result<(), TypeError> {
1320        let regex = crate::traits::get_media_regex();
1321        let mut new_content: Vec<ContentBlockParam> = Vec::with_capacity(self.content.len());
1322        for part in self.content.drain(..) {
1323            match part.inner {
1324                ContentBlock::Text(t) if regex.is_match(&t.text) => {
1325                    let cache = t.cache_control.clone();
1326                    let citations = t.citations.clone();
1327                    let pieces = crate::traits::split_text_on_media(&t.text, regex);
1328                    let mut first = true;
1329                    for s in pieces {
1330                        if s.is_empty() {
1331                            continue;
1332                        }
1333                        let block = TextBlockParam::new_rs(
1334                            s,
1335                            if first { cache.clone() } else { None },
1336                            if first { citations.clone() } else { None },
1337                        );
1338                        new_content.push(ContentBlockParam {
1339                            inner: ContentBlock::Text(block),
1340                        });
1341                        first = false;
1342                    }
1343                }
1344                other => new_content.push(ContentBlockParam { inner: other }),
1345            }
1346        }
1347        self.content = new_content;
1348        Ok(())
1349    }
1350
1351    fn bind_media_mut(
1352        &mut self,
1353        token: &str,
1354        media: &crate::prompt::media::MediaRef,
1355        _provider: &crate::Provider,
1356    ) -> Result<bool, TypeError> {
1357        use crate::prompt::media::{MediaKind, MediaSource};
1358
1359        let mut replaced = false;
1360        let mut new_content: Vec<ContentBlockParam> = Vec::with_capacity(self.content.len());
1361        for part in self.content.drain(..) {
1362            match &part.inner {
1363                ContentBlock::Text(t) if t.text.trim() == token => {
1364                    let new_inner = match (media.kind, &media.source) {
1365                        (MediaKind::Image, MediaSource::Url { url, .. }) => {
1366                            ContentBlock::Image(ImageBlockParam {
1367                                source: ImageSource::Url(UrlImageSource::new(url.clone())),
1368                                cache_control: None,
1369                                r#type: IMAGE_TYPE.to_string(),
1370                            })
1371                        }
1372                        (MediaKind::Image, MediaSource::Base64 { mime_type, data }) => {
1373                            ContentBlock::Image(ImageBlockParam {
1374                                source: ImageSource::Base64(Base64ImageSource {
1375                                    media_type: mime_type.clone(),
1376                                    data: data.clone(),
1377                                    r#type: BASE64_TYPE.to_string(),
1378                                }),
1379                                cache_control: None,
1380                                r#type: IMAGE_TYPE.to_string(),
1381                            })
1382                        }
1383                        (MediaKind::Document, MediaSource::Url { url, .. }) => {
1384                            ContentBlock::Document(DocumentBlockParam {
1385                                source: DocumentSource::Url(UrlPDFSource::new(url.clone())),
1386                                cache_control: None,
1387                                title: None,
1388                                context: None,
1389                                r#type: DOCUMENT_TYPE.to_string(),
1390                                citations: None,
1391                            })
1392                        }
1393                        (MediaKind::Document, MediaSource::Base64 { mime_type, data }) => {
1394                            ContentBlock::Document(DocumentBlockParam {
1395                                source: DocumentSource::Base64(Base64PDFSource {
1396                                    media_type: mime_type.clone(),
1397                                    data: data.clone(),
1398                                    r#type: BASE64_TYPE.to_string(),
1399                                }),
1400                                cache_control: None,
1401                                title: None,
1402                                context: None,
1403                                r#type: DOCUMENT_TYPE.to_string(),
1404                                citations: None,
1405                            })
1406                        }
1407                    };
1408                    new_content.push(ContentBlockParam { inner: new_inner });
1409                    replaced = true;
1410                }
1411                ContentBlock::Text(t) if t.text.contains(token) => {
1412                    return Err(TypeError::MediaPlaceholderNotIsolated {
1413                        name: crate::traits::extract_token_name(token),
1414                    });
1415                }
1416                _ => new_content.push(part),
1417            }
1418        }
1419        self.content = new_content;
1420        Ok(replaced)
1421    }
1422}
1423
1424impl MessageParam {
1425    /// Helper function to create a MessageParam from a single TextBlockParam
1426    pub fn to_text_block_param(&self) -> Result<TextBlockParam, TypeError> {
1427        if self.content.len() != 1 {
1428            return Err(TypeError::InvalidInput(
1429                "MessageParam must contain exactly one content block to convert to TextBlockParam"
1430                    .to_string(),
1431            ));
1432        }
1433
1434        match &self.content[0].inner {
1435            ContentBlock::Text(text_block) => Ok(text_block.clone()),
1436            _ => Err(TypeError::InvalidInput(
1437                "Content block is not of type TextBlockParam".to_string(),
1438            )),
1439        }
1440    }
1441}
1442
1443impl MessageFactory for MessageParam {
1444    fn from_text(content: String, role: &str) -> Result<Self, TypeError> {
1445        let text_block = TextBlockParam::new_rs(content, None, None);
1446        let content_block = ContentBlockParam {
1447            inner: ContentBlock::Text(text_block),
1448        };
1449
1450        Ok(Self {
1451            role: role.to_string(),
1452            content: vec![content_block],
1453        })
1454    }
1455}
1456
1457impl MessageConversion for MessageParam {
1458    fn to_anthropic_message(&self) -> Result<Self, TypeError> {
1459        // Currently, MessageParam is already in the Anthropic format
1460        Err(TypeError::CantConvertSelf)
1461    }
1462
1463    fn to_google_message(
1464        &self,
1465    ) -> Result<crate::google::v1::generate::request::GeminiContent, TypeError> {
1466        // Extract text content from all text blocks
1467        let mut parts = Vec::new();
1468
1469        for content_block in &self.content {
1470            match &content_block.inner {
1471                ContentBlock::Text(text_block) => {
1472                    parts.push(Part {
1473                        data: DataNum::Text(text_block.text.clone()),
1474                        thought: None,
1475                        thought_signature: None,
1476                        part_metadata: None,
1477                        media_resolution: None,
1478                        video_metadata: None,
1479                    });
1480                }
1481                _ => {
1482                    return Err(TypeError::UnsupportedConversion(
1483                        "Only text content blocks are currently supported for conversion"
1484                            .to_string(),
1485                    ));
1486                }
1487            }
1488        }
1489
1490        if parts.is_empty() {
1491            return Err(TypeError::UnsupportedConversion(
1492                "Message contains no text content to convert".to_string(),
1493            ));
1494        }
1495
1496        Ok(GeminiContent {
1497            role: self.role.clone(),
1498            parts,
1499        })
1500    }
1501
1502    fn to_openai_message(
1503        &self,
1504    ) -> Result<crate::openai::v1::chat::request::ChatMessage, TypeError> {
1505        // Extract text content from all text blocks
1506        let mut content_parts = Vec::new();
1507
1508        for content_block in &self.content {
1509            match &content_block.inner {
1510                ContentBlock::Text(text_block) => {
1511                    content_parts.push(ContentPart::Text(TextContentPart::new(
1512                        text_block.text.clone(),
1513                    )));
1514                }
1515                _ => {
1516                    return Err(TypeError::UnsupportedConversion(
1517                        "Only text content blocks are currently supported for conversion"
1518                            .to_string(),
1519                    ));
1520                }
1521            }
1522        }
1523
1524        if content_parts.is_empty() {
1525            return Err(TypeError::UnsupportedConversion(
1526                "Message contains no text content to convert".to_string(),
1527            ));
1528        }
1529
1530        Ok(ChatMessage {
1531            role: self.role.clone(),
1532            content: content_parts,
1533            name: None,
1534        })
1535    }
1536}
1537
1538#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1539#[pyclass(from_py_object)]
1540pub struct Metadata {
1541    #[serde(skip_serializing_if = "Option::is_none")]
1542    pub user_id: Option<String>,
1543}
1544
1545#[pymethods]
1546impl Metadata {
1547    #[new]
1548    #[pyo3(signature = (user_id=None))]
1549    pub fn new(user_id: Option<String>) -> Self {
1550        Self { user_id }
1551    }
1552}
1553
1554#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1555#[pyclass(from_py_object)]
1556pub struct CacheControl {
1557    #[serde(rename = "type")]
1558    pub cache_type: String, // "ephemeral"
1559    #[serde(skip_serializing_if = "Option::is_none")]
1560    pub ttl: Option<String>, // "5m" or "1h"
1561}
1562
1563#[pymethods]
1564impl CacheControl {
1565    #[new]
1566    #[pyo3(signature = (cache_type, ttl=None))]
1567    pub fn new(cache_type: String, ttl: Option<String>) -> Self {
1568        Self { cache_type, ttl }
1569    }
1570}
1571
1572#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1573#[pyclass(from_py_object, name = "AnthropicTool")]
1574pub struct Tool {
1575    pub name: String,
1576    #[serde(skip_serializing_if = "Option::is_none")]
1577    pub description: Option<String>,
1578    pub input_schema: Value,
1579    #[serde(skip_serializing_if = "Option::is_none")]
1580    pub cache_control: Option<CacheControl>,
1581}
1582
1583#[pymethods]
1584impl Tool {
1585    #[new]
1586    #[pyo3(signature = (
1587        name,
1588        input_schema,
1589        description=None,
1590        cache_control=None
1591    ))]
1592    pub fn new(
1593        name: String,
1594        input_schema: &Bound<'_, PyAny>,
1595        description: Option<String>,
1596        cache_control: Option<CacheControl>,
1597    ) -> Result<Self, UtilError> {
1598        Ok(Self {
1599            name,
1600            description,
1601            input_schema: depythonize(input_schema)?,
1602            cache_control,
1603        })
1604    }
1605}
1606
1607impl Tool {
1608    pub fn from_tool_agent_tool_definition(tool: &AgentToolDefinition) -> Result<Self, TypeError> {
1609        Ok(Self {
1610            name: tool.name.clone(),
1611            description: Some(tool.description.clone()),
1612            input_schema: tool.parameters.clone(),
1613            cache_control: None,
1614        })
1615    }
1616}
1617
1618#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1619#[pyclass(from_py_object, name = "AnthropicThinkingConfig")]
1620pub struct ThinkingConfig {
1621    #[pyo3(get)]
1622    pub r#type: String,
1623
1624    #[serde(skip_serializing_if = "Option::is_none")]
1625    #[pyo3(get)]
1626    pub budget_tokens: Option<i32>,
1627}
1628
1629#[pymethods]
1630impl ThinkingConfig {
1631    #[new]
1632    #[pyo3(signature = (r#type, budget_tokens=None))]
1633    pub fn new(r#type: String, budget_tokens: Option<i32>) -> Self {
1634        Self {
1635            r#type,
1636            budget_tokens,
1637        }
1638    }
1639}
1640
1641#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1642#[pyclass(from_py_object, name = "AnthropicToolChoice")]
1643pub struct ToolChoice {
1644    #[pyo3(get)]
1645    pub r#type: String, // "auto", "any", "tool", "none"
1646
1647    #[serde(skip_serializing_if = "Option::is_none")]
1648    #[pyo3(get)]
1649    disable_parallel_tool_use: Option<bool>,
1650
1651    #[serde(skip_serializing_if = "Option::is_none")]
1652    #[pyo3(get)]
1653    pub name: Option<String>,
1654}
1655
1656#[pymethods]
1657impl ToolChoice {
1658    #[new]
1659    #[pyo3(signature = (r#type, disable_parallel_tool_use=None, name=None))]
1660    pub fn new(
1661        r#type: String,
1662        disable_parallel_tool_use: Option<bool>,
1663        name: Option<String>,
1664    ) -> Result<Self, UtilError> {
1665        match name {
1666            Some(_) if r#type != "tool" => {
1667                return Err(UtilError::PyError(
1668                    "ToolChoice name can only be set if type is 'tool'".to_string(),
1669                ))
1670            }
1671            None if r#type == "tool" => {
1672                return Err(UtilError::PyError(
1673                    "ToolChoice of type 'tool' requires a name".to_string(),
1674                ))
1675            }
1676            _ => {}
1677        }
1678
1679        Ok(Self {
1680            r#type,
1681            disable_parallel_tool_use,
1682            name,
1683        })
1684    }
1685}
1686
1687#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1688#[pyclass(from_py_object)]
1689#[serde(default)]
1690pub struct AnthropicSettings {
1691    #[pyo3(get)]
1692    pub max_tokens: i32,
1693
1694    #[serde(skip_serializing_if = "Option::is_none")]
1695    #[pyo3(get)]
1696    pub metadata: Option<Metadata>,
1697
1698    #[serde(skip_serializing_if = "Option::is_none")]
1699    #[pyo3(get)]
1700    pub service_tier: Option<String>,
1701
1702    #[serde(skip_serializing_if = "Option::is_none")]
1703    #[pyo3(get)]
1704    pub stop_sequences: Option<Vec<String>>,
1705
1706    #[serde(skip_serializing_if = "Option::is_none")]
1707    #[pyo3(get)]
1708    pub stream: Option<bool>,
1709
1710    #[serde(skip_serializing_if = "Option::is_none")]
1711    #[pyo3(get)]
1712    pub system: Option<String>,
1713
1714    #[serde(skip_serializing_if = "Option::is_none")]
1715    #[pyo3(get)]
1716    pub temperature: Option<f32>,
1717
1718    #[serde(skip_serializing_if = "Option::is_none")]
1719    #[pyo3(get)]
1720    pub thinking: Option<ThinkingConfig>,
1721
1722    #[serde(skip_serializing_if = "Option::is_none")]
1723    #[pyo3(get)]
1724    pub tool_choice: Option<ToolChoice>,
1725
1726    #[serde(skip_serializing_if = "Option::is_none")]
1727    #[pyo3(get)]
1728    pub tools: Option<Vec<Tool>>,
1729
1730    #[serde(skip_serializing_if = "Option::is_none")]
1731    #[pyo3(get)]
1732    pub top_k: Option<i32>,
1733
1734    #[serde(skip_serializing_if = "Option::is_none")]
1735    #[pyo3(get)]
1736    pub top_p: Option<f32>,
1737
1738    #[serde(skip_serializing_if = "Option::is_none")]
1739    pub extra_body: Option<Value>,
1740}
1741
1742impl Default for AnthropicSettings {
1743    fn default() -> Self {
1744        Self {
1745            max_tokens: 4096,
1746            metadata: None,
1747            service_tier: None,
1748            stop_sequences: None,
1749            stream: Some(false),
1750            system: None,
1751            temperature: None,
1752            thinking: None,
1753            top_k: None,
1754            top_p: None,
1755            tools: None,
1756            tool_choice: None,
1757            extra_body: None,
1758        }
1759    }
1760}
1761
1762#[pymethods]
1763impl AnthropicSettings {
1764    #[new]
1765    #[pyo3(signature = (
1766        max_tokens=4096,
1767        metadata=None,
1768        service_tier=None,
1769        stop_sequences=None,
1770        stream=None,
1771        system =None,
1772        temperature=None,
1773        thinking=None,
1774        top_k=None,
1775        top_p=None,
1776        tools=None,
1777        tool_choice=None,
1778        extra_body=None
1779    ))]
1780    #[allow(clippy::too_many_arguments)]
1781    pub fn new(
1782        max_tokens: i32,
1783        metadata: Option<Metadata>,
1784        service_tier: Option<String>,
1785        stop_sequences: Option<Vec<String>>,
1786        stream: Option<bool>,
1787        system: Option<String>,
1788        temperature: Option<f32>,
1789        thinking: Option<ThinkingConfig>,
1790        top_k: Option<i32>,
1791        top_p: Option<f32>,
1792        tools: Option<Vec<Tool>>,
1793        tool_choice: Option<ToolChoice>,
1794        extra_body: Option<&Bound<'_, PyAny>>,
1795    ) -> Result<Self, UtilError> {
1796        let extra = match extra_body {
1797            Some(obj) => Some(depythonize(obj)?),
1798            None => None,
1799        };
1800
1801        Ok(Self {
1802            max_tokens,
1803            metadata,
1804            service_tier,
1805            stop_sequences,
1806            stream,
1807            system,
1808            temperature,
1809            thinking,
1810            top_k,
1811            top_p,
1812            tools,
1813            tool_choice,
1814            extra_body: extra,
1815        })
1816    }
1817
1818    pub fn __str__(&self) -> String {
1819        PyHelperFuncs::__str__(self)
1820    }
1821
1822    pub fn model_dump<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
1823        // iterate over each field in model_settings and add to the dict if it is not None
1824        let json = serde_json::to_value(self)?;
1825        Ok(pythonize(py, &json)?)
1826    }
1827
1828    pub fn settings_type(&self) -> SettingsType {
1829        SettingsType::Anthropic
1830    }
1831}
1832
1833impl AnthropicSettings {
1834    pub fn add_tools(&mut self, tools: Vec<AgentToolDefinition>) -> Result<(), TypeError> {
1835        let current_tools = self.tools.get_or_insert_with(Vec::new);
1836
1837        for tool in tools {
1838            let tool_param = Tool::from_tool_agent_tool_definition(&tool)?;
1839            current_tools.push(tool_param);
1840        }
1841
1842        Ok(())
1843    }
1844}
1845
1846#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1847pub struct AnthropicMessageRequestV1 {
1848    pub model: String,
1849    pub messages: Vec<MessageNum>,
1850    pub system: Vec<MessageNum>,
1851    #[serde(flatten)]
1852    pub settings: AnthropicSettings,
1853    #[serde(skip_serializing_if = "Option::is_none")]
1854    pub output_format: Option<Value>,
1855}
1856
1857pub(crate) fn create_structured_output_schema(json_schema: &Value) -> Value {
1858    serde_json::json!({
1859        "type": "json_schema",
1860        "schema": json_schema,
1861
1862    })
1863}
1864
1865impl RequestAdapter for AnthropicMessageRequestV1 {
1866    fn messages_mut(&mut self) -> &mut Vec<MessageNum> {
1867        &mut self.messages
1868    }
1869    fn messages(&self) -> &[MessageNum] {
1870        &self.messages
1871    }
1872    fn system_instructions(&self) -> Vec<&MessageNum> {
1873        self.system.iter().collect()
1874    }
1875    fn response_json_schema(&self) -> Option<&Value> {
1876        let schema = self.output_format.as_ref();
1877
1878        // if Some, get "schema" field
1879        // output_format: {
1880        //     "type": "json_schema",
1881        //     "schema": { ... }
1882        // }
1883        schema.and_then(|of| of.get("schema"))
1884    }
1885
1886    fn preprend_system_instructions(&mut self, messages: Vec<MessageNum>) -> Result<(), TypeError> {
1887        let mut combined = messages;
1888        combined.append(&mut self.system);
1889        self.system = combined;
1890        Ok(())
1891    }
1892
1893    fn get_py_system_instructions<'py>(
1894        &self,
1895        py: Python<'py>,
1896    ) -> Result<Bound<'py, PyList>, TypeError> {
1897        let py_system_instructions = PyList::empty(py);
1898        for system_msg in &self.system {
1899            py_system_instructions.append(system_msg.to_bound_py_object(py)?)?;
1900        }
1901
1902        Ok(py_system_instructions)
1903    }
1904
1905    fn model_settings<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
1906        let settings = self.settings.clone();
1907        Ok(settings.into_bound_py_any(py)?)
1908    }
1909
1910    fn to_request_body(&self) -> Result<Value, TypeError> {
1911        Ok(serde_json::to_value(self)?)
1912    }
1913    fn match_provider(&self, provider: &Provider) -> bool {
1914        *provider == Provider::Anthropic
1915    }
1916    fn build_provider_enum(
1917        messages: Vec<MessageNum>,
1918        system_instructions: Vec<MessageNum>,
1919        model: String,
1920        settings: ModelSettings,
1921        response_json_schema: Option<Value>,
1922    ) -> Result<ProviderRequest, TypeError> {
1923        let anthropic_settings = match settings {
1924            ModelSettings::AnthropicChat(s) => s,
1925            _ => AnthropicSettings::default(),
1926        };
1927
1928        let output_format =
1929            response_json_schema.map(|json_schema| create_structured_output_schema(&json_schema));
1930
1931        Ok(ProviderRequest::AnthropicV1(AnthropicMessageRequestV1 {
1932            model,
1933            messages,
1934            system: system_instructions,
1935            settings: anthropic_settings,
1936            output_format,
1937        }))
1938    }
1939
1940    fn set_response_json_schema(&mut self, response_json_schema: Option<Value>) {
1941        self.output_format =
1942            response_json_schema.map(|json_schema| create_structured_output_schema(&json_schema));
1943    }
1944
1945    fn add_tools(&mut self, tools: Vec<AgentToolDefinition>) -> Result<(), TypeError> {
1946        self.settings.add_tools(tools)
1947    }
1948}
1949
1950#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1951#[pyclass(from_py_object)]
1952pub struct SystemPrompt {
1953    #[pyo3(get)]
1954    #[serde(flatten)]
1955    pub content: Vec<TextBlockParam>,
1956}
1957
1958#[pymethods]
1959impl SystemPrompt {
1960    /// Create a new SystemPrompt
1961    /// Accepts either a single string or a list of TextBlockParams
1962    /// # Arguments
1963    /// * `content` - Either a string or a list of TextBlockParams
1964    /// # Returns
1965    /// * `SystemPrompt` - The created SystemPrompt
1966    /// Errors
1967    /// * `TypeError` - If the content is not a string or a list of TextBlockParams
1968    #[new]
1969    pub fn new(content: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
1970        let content_blocks: Vec<TextBlockParam> =
1971            if content.is_instance_of::<pyo3::types::PyString>() {
1972                let text = content.extract::<String>()?;
1973                let text_block = TextBlockParam::new(text, None, None)?;
1974                vec![text_block]
1975            } else if content.is_instance_of::<pyo3::types::PyList>() {
1976                let content_list = content.extract::<Vec<Bound<'_, PyAny>>>()?;
1977                let mut blocks = Vec::new();
1978                for item in content_list {
1979                    let text_block = item.extract::<TextBlockParam>().map_err(|_| {
1980                        TypeError::InvalidInput(
1981                            "All items in the list must be TextBlockParam".to_string(),
1982                        )
1983                    })?;
1984                    blocks.push(text_block);
1985                }
1986                blocks
1987            } else {
1988                return Err(TypeError::InvalidInput(
1989                    "Content must be either a string or a list of TextBlockParam".to_string(),
1990                ));
1991            };
1992
1993        Ok(Self {
1994            content: content_blocks,
1995        })
1996    }
1997}
potato_type/anthropic/v1/request.rs

potato_type/anthropic/v1/
request.rs
