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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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]
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 placeholder = format!("${{{}}}", name);
1246
1247        for part in &mut self.content {
1248            if let ContentBlock::Text(text_part) = &mut part.inner {
1249                text_part.text = text_part.text.replace(&placeholder, value);
1250            }
1251        }
1252
1253        Ok(())
1254    }
1255
1256    fn bind(&self, name: &str, value: &str) -> Result<Self, TypeError>
1257    where
1258        Self: Sized,
1259    {
1260        let mut new_message = self.clone();
1261        new_message.bind_mut(name, value)?;
1262        Ok(new_message)
1263    }
1264
1265    fn extract_variables(&self) -> Vec<String> {
1266        let mut variables = HashSet::new();
1267
1268        // Lazily initialize regex to avoid recompilation
1269        let regex = get_var_regex();
1270
1271        // Extract variables from all text content parts
1272        for part in &self.content {
1273            if let ContentBlock::Text(text_part) = &part.inner {
1274                for captures in regex.captures_iter(&text_part.text) {
1275                    if let Some(name) = captures.get(1) {
1276                        variables.insert(name.as_str().to_string());
1277                    }
1278                }
1279            }
1280        }
1281
1282        // Convert HashSet to Vec for return
1283        variables.into_iter().collect()
1284    }
1285
1286    fn from_text(content: String, role: &str) -> Result<Self, TypeError> {
1287        Ok(Self {
1288            role: role.to_string(),
1289            content: vec![ContentBlockParam {
1290                inner: ContentBlock::Text(TextBlockParam::new_rs(content, None, None)),
1291            }],
1292        })
1293    }
1294}
1295
1296impl MessageParam {
1297    /// Helper function to create a MessageParam from a single TextBlockParam
1298    pub fn to_text_block_param(&self) -> Result<TextBlockParam, TypeError> {
1299        if self.content.len() != 1 {
1300            return Err(TypeError::InvalidInput(
1301                "MessageParam must contain exactly one content block to convert to TextBlockParam"
1302                    .to_string(),
1303            ));
1304        }
1305
1306        match &self.content[0].inner {
1307            ContentBlock::Text(text_block) => Ok(text_block.clone()),
1308            _ => Err(TypeError::InvalidInput(
1309                "Content block is not of type TextBlockParam".to_string(),
1310            )),
1311        }
1312    }
1313}
1314
1315impl MessageFactory for MessageParam {
1316    fn from_text(content: String, role: &str) -> Result<Self, TypeError> {
1317        let text_block = TextBlockParam::new_rs(content, None, None);
1318        let content_block = ContentBlockParam {
1319            inner: ContentBlock::Text(text_block),
1320        };
1321
1322        Ok(Self {
1323            role: role.to_string(),
1324            content: vec![content_block],
1325        })
1326    }
1327}
1328
1329impl MessageConversion for MessageParam {
1330    fn to_anthropic_message(&self) -> Result<Self, TypeError> {
1331        // Currently, MessageParam is already in the Anthropic format
1332        Err(TypeError::CantConvertSelf)
1333    }
1334
1335    fn to_google_message(
1336        &self,
1337    ) -> Result<crate::google::v1::generate::request::GeminiContent, TypeError> {
1338        // Extract text content from all text blocks
1339        let mut parts = Vec::new();
1340
1341        for content_block in &self.content {
1342            match &content_block.inner {
1343                ContentBlock::Text(text_block) => {
1344                    parts.push(Part {
1345                        data: DataNum::Text(text_block.text.clone()),
1346                        thought: None,
1347                        thought_signature: None,
1348                        part_metadata: None,
1349                        media_resolution: None,
1350                        video_metadata: None,
1351                    });
1352                }
1353                _ => {
1354                    return Err(TypeError::UnsupportedConversion(
1355                        "Only text content blocks are currently supported for conversion"
1356                            .to_string(),
1357                    ));
1358                }
1359            }
1360        }
1361
1362        if parts.is_empty() {
1363            return Err(TypeError::UnsupportedConversion(
1364                "Message contains no text content to convert".to_string(),
1365            ));
1366        }
1367
1368        Ok(GeminiContent {
1369            role: self.role.clone(),
1370            parts,
1371        })
1372    }
1373
1374    fn to_openai_message(
1375        &self,
1376    ) -> Result<crate::openai::v1::chat::request::ChatMessage, TypeError> {
1377        // Extract text content from all text blocks
1378        let mut content_parts = Vec::new();
1379
1380        for content_block in &self.content {
1381            match &content_block.inner {
1382                ContentBlock::Text(text_block) => {
1383                    content_parts.push(ContentPart::Text(TextContentPart::new(
1384                        text_block.text.clone(),
1385                    )));
1386                }
1387                _ => {
1388                    return Err(TypeError::UnsupportedConversion(
1389                        "Only text content blocks are currently supported for conversion"
1390                            .to_string(),
1391                    ));
1392                }
1393            }
1394        }
1395
1396        if content_parts.is_empty() {
1397            return Err(TypeError::UnsupportedConversion(
1398                "Message contains no text content to convert".to_string(),
1399            ));
1400        }
1401
1402        Ok(ChatMessage {
1403            role: self.role.clone(),
1404            content: content_parts,
1405            name: None,
1406        })
1407    }
1408}
1409
1410#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1411#[pyclass]
1412pub struct Metadata {
1413    #[serde(skip_serializing_if = "Option::is_none")]
1414    pub user_id: Option<String>,
1415}
1416
1417#[pymethods]
1418impl Metadata {
1419    #[new]
1420    #[pyo3(signature = (user_id=None))]
1421    pub fn new(user_id: Option<String>) -> Self {
1422        Self { user_id }
1423    }
1424}
1425
1426#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1427#[pyclass]
1428pub struct CacheControl {
1429    #[serde(rename = "type")]
1430    pub cache_type: String, // "ephemeral"
1431    #[serde(skip_serializing_if = "Option::is_none")]
1432    pub ttl: Option<String>, // "5m" or "1h"
1433}
1434
1435#[pymethods]
1436impl CacheControl {
1437    #[new]
1438    #[pyo3(signature = (cache_type, ttl=None))]
1439    pub fn new(cache_type: String, ttl: Option<String>) -> Self {
1440        Self { cache_type, ttl }
1441    }
1442}
1443
1444#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1445#[pyclass(name = "AnthropicTool")]
1446pub struct Tool {
1447    pub name: String,
1448    #[serde(skip_serializing_if = "Option::is_none")]
1449    pub description: Option<String>,
1450    pub input_schema: Value,
1451    #[serde(skip_serializing_if = "Option::is_none")]
1452    pub cache_control: Option<CacheControl>,
1453}
1454
1455#[pymethods]
1456impl Tool {
1457    #[new]
1458    #[pyo3(signature = (
1459        name,
1460        input_schema,
1461        description=None,
1462        cache_control=None
1463    ))]
1464    pub fn new(
1465        name: String,
1466        input_schema: &Bound<'_, PyAny>,
1467        description: Option<String>,
1468        cache_control: Option<CacheControl>,
1469    ) -> Result<Self, UtilError> {
1470        Ok(Self {
1471            name,
1472            description,
1473            input_schema: depythonize(input_schema)?,
1474            cache_control,
1475        })
1476    }
1477}
1478
1479impl Tool {
1480    pub fn from_tool_agent_tool_definition(tool: &AgentToolDefinition) -> Result<Self, TypeError> {
1481        Ok(Self {
1482            name: tool.name.clone(),
1483            description: Some(tool.description.clone()),
1484            input_schema: tool.parameters.clone(),
1485            cache_control: None,
1486        })
1487    }
1488}
1489
1490#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1491#[pyclass(name = "AnthropicThinkingConfig")]
1492pub struct ThinkingConfig {
1493    #[pyo3(get)]
1494    pub r#type: String,
1495
1496    #[serde(skip_serializing_if = "Option::is_none")]
1497    #[pyo3(get)]
1498    pub budget_tokens: Option<i32>,
1499}
1500
1501#[pymethods]
1502impl ThinkingConfig {
1503    #[new]
1504    #[pyo3(signature = (r#type, budget_tokens=None))]
1505    pub fn new(r#type: String, budget_tokens: Option<i32>) -> Self {
1506        Self {
1507            r#type,
1508            budget_tokens,
1509        }
1510    }
1511}
1512
1513#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1514#[pyclass(name = "AnthropicToolChoice")]
1515pub struct ToolChoice {
1516    #[pyo3(get)]
1517    pub r#type: String, // "auto", "any", "tool", "none"
1518
1519    #[serde(skip_serializing_if = "Option::is_none")]
1520    #[pyo3(get)]
1521    disable_parallel_tool_use: Option<bool>,
1522
1523    #[serde(skip_serializing_if = "Option::is_none")]
1524    #[pyo3(get)]
1525    pub name: Option<String>,
1526}
1527
1528#[pymethods]
1529impl ToolChoice {
1530    #[new]
1531    #[pyo3(signature = (r#type, disable_parallel_tool_use=None, name=None))]
1532    pub fn new(
1533        r#type: String,
1534        disable_parallel_tool_use: Option<bool>,
1535        name: Option<String>,
1536    ) -> Result<Self, UtilError> {
1537        match name {
1538            Some(_) if r#type != "tool" => {
1539                return Err(UtilError::PyError(
1540                    "ToolChoice name can only be set if type is 'tool'".to_string(),
1541                ))
1542            }
1543            None if r#type == "tool" => {
1544                return Err(UtilError::PyError(
1545                    "ToolChoice of type 'tool' requires a name".to_string(),
1546                ))
1547            }
1548            _ => {}
1549        }
1550
1551        Ok(Self {
1552            r#type,
1553            disable_parallel_tool_use,
1554            name,
1555        })
1556    }
1557}
1558
1559#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1560#[pyclass]
1561#[serde(default)]
1562pub struct AnthropicSettings {
1563    #[pyo3(get)]
1564    pub max_tokens: i32,
1565
1566    #[serde(skip_serializing_if = "Option::is_none")]
1567    #[pyo3(get)]
1568    pub metadata: Option<Metadata>,
1569
1570    #[serde(skip_serializing_if = "Option::is_none")]
1571    #[pyo3(get)]
1572    pub service_tier: Option<String>,
1573
1574    #[serde(skip_serializing_if = "Option::is_none")]
1575    #[pyo3(get)]
1576    pub stop_sequences: Option<Vec<String>>,
1577
1578    #[serde(skip_serializing_if = "Option::is_none")]
1579    #[pyo3(get)]
1580    pub stream: Option<bool>,
1581
1582    #[serde(skip_serializing_if = "Option::is_none")]
1583    #[pyo3(get)]
1584    pub system: Option<String>,
1585
1586    #[serde(skip_serializing_if = "Option::is_none")]
1587    #[pyo3(get)]
1588    pub temperature: Option<f32>,
1589
1590    #[serde(skip_serializing_if = "Option::is_none")]
1591    #[pyo3(get)]
1592    pub thinking: Option<ThinkingConfig>,
1593
1594    #[serde(skip_serializing_if = "Option::is_none")]
1595    #[pyo3(get)]
1596    pub tool_choice: Option<ToolChoice>,
1597
1598    #[serde(skip_serializing_if = "Option::is_none")]
1599    #[pyo3(get)]
1600    pub tools: Option<Vec<Tool>>,
1601
1602    #[serde(skip_serializing_if = "Option::is_none")]
1603    #[pyo3(get)]
1604    pub top_k: Option<i32>,
1605
1606    #[serde(skip_serializing_if = "Option::is_none")]
1607    #[pyo3(get)]
1608    pub top_p: Option<f32>,
1609
1610    #[serde(skip_serializing_if = "Option::is_none")]
1611    pub extra_body: Option<Value>,
1612}
1613
1614impl Default for AnthropicSettings {
1615    fn default() -> Self {
1616        Self {
1617            max_tokens: 4096,
1618            metadata: None,
1619            service_tier: None,
1620            stop_sequences: None,
1621            stream: Some(false),
1622            system: None,
1623            temperature: None,
1624            thinking: None,
1625            top_k: None,
1626            top_p: None,
1627            tools: None,
1628            tool_choice: None,
1629            extra_body: None,
1630        }
1631    }
1632}
1633
1634#[pymethods]
1635impl AnthropicSettings {
1636    #[new]
1637    #[pyo3(signature = (
1638        max_tokens=4096,
1639        metadata=None,
1640        service_tier=None,
1641        stop_sequences=None,
1642        stream=None,
1643        system =None,
1644        temperature=None,
1645        thinking=None,
1646        top_k=None,
1647        top_p=None,
1648        tools=None,
1649        tool_choice=None,
1650        extra_body=None
1651    ))]
1652    #[allow(clippy::too_many_arguments)]
1653    pub fn new(
1654        max_tokens: i32,
1655        metadata: Option<Metadata>,
1656        service_tier: Option<String>,
1657        stop_sequences: Option<Vec<String>>,
1658        stream: Option<bool>,
1659        system: Option<String>,
1660        temperature: Option<f32>,
1661        thinking: Option<ThinkingConfig>,
1662        top_k: Option<i32>,
1663        top_p: Option<f32>,
1664        tools: Option<Vec<Tool>>,
1665        tool_choice: Option<ToolChoice>,
1666        extra_body: Option<&Bound<'_, PyAny>>,
1667    ) -> Result<Self, UtilError> {
1668        let extra = match extra_body {
1669            Some(obj) => Some(depythonize(obj)?),
1670            None => None,
1671        };
1672
1673        Ok(Self {
1674            max_tokens,
1675            metadata,
1676            service_tier,
1677            stop_sequences,
1678            stream,
1679            system,
1680            temperature,
1681            thinking,
1682            top_k,
1683            top_p,
1684            tools,
1685            tool_choice,
1686            extra_body: extra,
1687        })
1688    }
1689
1690    pub fn __str__(&self) -> String {
1691        PyHelperFuncs::__str__(self)
1692    }
1693
1694    pub fn model_dump<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
1695        // iterate over each field in model_settings and add to the dict if it is not None
1696        let json = serde_json::to_value(self)?;
1697        Ok(pythonize(py, &json)?)
1698    }
1699
1700    pub fn settings_type(&self) -> SettingsType {
1701        SettingsType::Anthropic
1702    }
1703}
1704
1705impl AnthropicSettings {
1706    pub fn add_tools(&mut self, tools: Vec<AgentToolDefinition>) -> Result<(), TypeError> {
1707        let current_tools = self.tools.get_or_insert_with(Vec::new);
1708
1709        for tool in tools {
1710            let tool_param = Tool::from_tool_agent_tool_definition(&tool)?;
1711            current_tools.push(tool_param);
1712        }
1713
1714        Ok(())
1715    }
1716}
1717
1718#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1719pub struct AnthropicMessageRequestV1 {
1720    pub model: String,
1721    pub messages: Vec<MessageNum>,
1722    pub system: Vec<MessageNum>,
1723    #[serde(flatten)]
1724    pub settings: AnthropicSettings,
1725    #[serde(skip_serializing_if = "Option::is_none")]
1726    pub output_format: Option<Value>,
1727}
1728
1729pub(crate) fn create_structured_output_schema(json_schema: &Value) -> Value {
1730    serde_json::json!({
1731        "type": "json_schema",
1732        "schema": json_schema,
1733
1734    })
1735}
1736
1737impl RequestAdapter for AnthropicMessageRequestV1 {
1738    fn messages_mut(&mut self) -> &mut Vec<MessageNum> {
1739        &mut self.messages
1740    }
1741    fn messages(&self) -> &[MessageNum] {
1742        &self.messages
1743    }
1744    fn system_instructions(&self) -> Vec<&MessageNum> {
1745        self.system.iter().collect()
1746    }
1747    fn response_json_schema(&self) -> Option<&Value> {
1748        let schema = self.output_format.as_ref();
1749
1750        // if Some, get "schema" field
1751        // output_format: {
1752        //     "type": "json_schema",
1753        //     "schema": { ... }
1754        // }
1755        schema.and_then(|of| of.get("schema"))
1756    }
1757
1758    fn preprend_system_instructions(&mut self, messages: Vec<MessageNum>) -> Result<(), TypeError> {
1759        let mut combined = messages;
1760        combined.append(&mut self.system);
1761        self.system = combined;
1762        Ok(())
1763    }
1764
1765    fn get_py_system_instructions<'py>(
1766        &self,
1767        py: Python<'py>,
1768    ) -> Result<Bound<'py, PyList>, TypeError> {
1769        let py_system_instructions = PyList::empty(py);
1770        for system_msg in &self.system {
1771            py_system_instructions.append(system_msg.to_bound_py_object(py)?)?;
1772        }
1773
1774        Ok(py_system_instructions)
1775    }
1776
1777    fn model_settings<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
1778        let settings = self.settings.clone();
1779        Ok(settings.into_bound_py_any(py)?)
1780    }
1781
1782    fn to_request_body(&self) -> Result<Value, TypeError> {
1783        Ok(serde_json::to_value(self)?)
1784    }
1785    fn match_provider(&self, provider: &Provider) -> bool {
1786        *provider == Provider::Anthropic
1787    }
1788    fn build_provider_enum(
1789        messages: Vec<MessageNum>,
1790        system_instructions: Vec<MessageNum>,
1791        model: String,
1792        settings: ModelSettings,
1793        response_json_schema: Option<Value>,
1794    ) -> Result<ProviderRequest, TypeError> {
1795        let anthropic_settings = match settings {
1796            ModelSettings::AnthropicChat(s) => s,
1797            _ => AnthropicSettings::default(),
1798        };
1799
1800        let output_format =
1801            response_json_schema.map(|json_schema| create_structured_output_schema(&json_schema));
1802
1803        Ok(ProviderRequest::AnthropicV1(AnthropicMessageRequestV1 {
1804            model,
1805            messages,
1806            system: system_instructions,
1807            settings: anthropic_settings,
1808            output_format,
1809        }))
1810    }
1811
1812    fn set_response_json_schema(&mut self, response_json_schema: Option<Value>) {
1813        self.output_format =
1814            response_json_schema.map(|json_schema| create_structured_output_schema(&json_schema));
1815    }
1816
1817    fn add_tools(&mut self, tools: Vec<AgentToolDefinition>) -> Result<(), TypeError> {
1818        self.settings.add_tools(tools)
1819    }
1820}
1821
1822#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
1823#[pyclass]
1824pub struct SystemPrompt {
1825    #[pyo3(get)]
1826    #[serde(flatten)]
1827    pub content: Vec<TextBlockParam>,
1828}
1829
1830#[pymethods]
1831impl SystemPrompt {
1832    /// Create a new SystemPrompt
1833    /// Accepts either a single string or a list of TextBlockParams
1834    /// # Arguments
1835    /// * `content` - Either a string or a list of TextBlockParams
1836    /// # Returns
1837    /// * `SystemPrompt` - The created SystemPrompt
1838    /// Errors
1839    /// * `TypeError` - If the content is not a string or a list of TextBlockParams
1840    #[new]
1841    pub fn new(content: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
1842        let content_blocks: Vec<TextBlockParam> =
1843            if content.is_instance_of::<pyo3::types::PyString>() {
1844                let text = content.extract::<String>()?;
1845                let text_block = TextBlockParam::new(text, None, None)?;
1846                vec![text_block]
1847            } else if content.is_instance_of::<pyo3::types::PyList>() {
1848                let content_list = content.extract::<Vec<Bound<'_, PyAny>>>()?;
1849                let mut blocks = Vec::new();
1850                for item in content_list {
1851                    let text_block = item.extract::<TextBlockParam>().map_err(|_| {
1852                        TypeError::InvalidInput(
1853                            "All items in the list must be TextBlockParam".to_string(),
1854                        )
1855                    })?;
1856                    blocks.push(text_block);
1857                }
1858                blocks
1859            } else {
1860                return Err(TypeError::InvalidInput(
1861                    "Content must be either a string or a list of TextBlockParam".to_string(),
1862                ));
1863            };
1864
1865        Ok(Self {
1866            content: content_blocks,
1867        })
1868    }
1869}
potato_type/anthropic/v1/request.rs

potato_type/anthropic/v1/
request.rs
