potato_type/prompt/
types.rs

1use crate::anthropic::v1::request::MessageParam as AnthropicMessage;
2use crate::anthropic::v1::request::TextBlockParam;
3use crate::anthropic::v1::response::ResponseContentBlock;
4use crate::google::v1::generate::Candidate;
5use crate::google::v1::generate::GeminiContent;
6use crate::google::PredictResponse;
7use crate::openai::v1::chat::request::ChatMessage as OpenAIChatMessage;
8use crate::openai::v1::Choice;
9use crate::traits::MessageConversion;
10use crate::traits::PromptMessageExt;
11use crate::Provider;
12use crate::{StructuredOutput, TypeError};
13use potato_util::PyHelperFuncs;
14use pyo3::types::PyAnyMethods;
15use pyo3::{prelude::*, IntoPyObjectExt};
16use pythonize::{depythonize, pythonize};
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use std::fmt::Display;
21use tracing::error;
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
23#[pyclass(from_py_object)]
24pub enum Role {
25    User,
26    Assistant,
27    Developer,
28    Tool,
29    Model,
30    System,
31}
32
33#[pymethods]
34impl Role {
35    #[pyo3(name = "as_str")]
36    pub fn as_str_py(&self) -> &'static str {
37        self.as_str()
38    }
39}
40
41impl Role {
42    /// Returns the string representation of the role
43    pub const fn as_str(&self) -> &'static str {
44        match self {
45            Role::User => "user",
46            Role::Assistant => "assistant",
47            Role::Developer => "developer",
48            Role::Tool => "tool",
49            Role::Model => "model",
50            Role::System => "system",
51        }
52    }
53}
54
55impl Display for Role {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        match self {
58            Role::User => write!(f, "user"),
59            Role::Assistant => write!(f, "assistant"),
60            Role::Developer => write!(f, "developer"),
61            Role::Tool => write!(f, "tool"),
62            Role::Model => write!(f, "model"),
63            Role::System => write!(f, "system"),
64        }
65    }
66}
67
68impl From<Role> for &str {
69    fn from(role: Role) -> Self {
70        match role {
71            Role::User => "user",
72            Role::Assistant => "assistant",
73            Role::Developer => "developer",
74            Role::Tool => "tool",
75            Role::Model => "model",
76            Role::System => "system",
77        }
78    }
79}
80
81pub trait DeserializePromptValExt: for<'de> serde::Deserialize<'de> {
82    /// Validates and deserializes a JSON value into its struct type.
83    ///
84    /// # Arguments
85    /// * `value` - The JSON value to deserialize
86    ///
87    /// # Returns
88    /// * `Result<Self, serde_json::Error>` - The deserialized value or error
89    fn model_validate_json(value: &Value) -> Result<Self, serde_json::Error> {
90        serde_json::from_value(value.clone())
91    }
92}
93
94pub fn get_pydantic_module<'py>(py: Python<'py>, module_name: &str) -> PyResult<Bound<'py, PyAny>> {
95    py.import("pydantic_ai")?.getattr(module_name)
96}
97
98/// Checks if an object is a subclass of a pydantic BaseModel. This is used when validating structured outputs
99/// # Arguments
100/// * `py` - The Python interpreter instance
101/// * `object` - The object to check
102/// # Returns
103/// A boolean indicating whether the object is a subclass of pydantic.BaseModel
104pub fn check_pydantic_model<'py>(
105    py: Python<'py>,
106    object: &Bound<'_, PyAny>,
107) -> Result<bool, TypeError> {
108    // check pydantic import. Return false if it fails
109    let pydantic = match py.import("pydantic").map_err(|e| {
110        error!("Failed to import pydantic: {}", e);
111        false
112    }) {
113        Ok(pydantic) => pydantic,
114        Err(_) => return Ok(false),
115    };
116
117    // get builtin subclass
118    let is_subclass = py.import("builtins")?.getattr("issubclass")?;
119
120    // Need to check if provided object is a basemodel
121    let basemodel = pydantic.getattr("BaseModel")?;
122    let matched = is_subclass.call1((object, basemodel))?.extract::<bool>()?;
123
124    Ok(matched)
125}
126
127/// Generate a JSON schema from a pydantic BaseModel object.
128/// # Arguments
129/// * `object` - The pydantic BaseModel object to generate the schema from.
130/// # Returns
131/// A JSON schema as a serde_json::Value.
132fn get_json_schema_from_basemodel(object: &Bound<'_, PyAny>) -> Result<Value, TypeError> {
133    // call staticmethod .model_json_schema()
134    let schema = object.getattr("model_json_schema")?.call1(())?;
135
136    let mut schema: Value = depythonize(&schema)?;
137
138    // ensure schema as additionalProperties set to false
139    if let Some(additional_properties) = schema.get_mut("additionalProperties") {
140        *additional_properties = serde_json::json!(false);
141    } else {
142        schema
143            .as_object_mut()
144            .unwrap()
145            .insert("additionalProperties".to_string(), serde_json::json!(false));
146    }
147
148    Ok(schema)
149}
150
151fn parse_pydantic_model<'py>(
152    py: Python<'py>,
153    object: &Bound<'_, PyAny>,
154) -> Result<Option<Value>, TypeError> {
155    let is_subclass = check_pydantic_model(py, object)?;
156    if is_subclass {
157        Ok(Some(get_json_schema_from_basemodel(object)?))
158    } else {
159        Ok(None)
160    }
161}
162
163pub fn check_response_type(object: &Bound<'_, PyAny>) -> Result<Option<ResponseType>, TypeError> {
164    // try calling staticmethod response_type()
165    let response_type = match object.getattr("response_type") {
166        Ok(method) => {
167            if method.is_callable() {
168                let response_type: ResponseType = method.call0()?.extract()?;
169                Some(response_type)
170            } else {
171                None
172            }
173        }
174        Err(_) => None,
175    };
176
177    Ok(response_type)
178}
179
180fn get_json_schema_from_response_type(response_type: &ResponseType) -> Result<Value, TypeError> {
181    match response_type {
182        ResponseType::Score => Ok(Score::get_structured_output_schema()),
183        _ => {
184            // If the response type is not recognized, return None
185            Err(TypeError::Error(format!(
186                "Unsupported response type: {response_type}"
187            )))
188        }
189    }
190}
191
192pub fn parse_response_to_json<'py>(
193    py: Python<'py>,
194    object: &Bound<'_, PyAny>,
195) -> Result<(ResponseType, Option<Value>), TypeError> {
196    // check if object is a pydantic model
197    let is_pydantic_model = check_pydantic_model(py, object)?;
198    if is_pydantic_model {
199        return Ok((ResponseType::Pydantic, parse_pydantic_model(py, object)?));
200    }
201
202    // check if object has response_type method
203    let response_type = check_response_type(object)?;
204    if let Some(response_type) = response_type {
205        return Ok((
206            response_type.clone(),
207            Some(get_json_schema_from_response_type(&response_type)?),
208        ));
209    }
210
211    Ok((ResponseType::Null, None))
212}
213
214#[pyclass(from_py_object)]
215#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
216#[serde(deny_unknown_fields)] // ensure strict validation
217pub struct Score {
218    #[pyo3(get)]
219    #[schemars(range(min = 1, max = 5))]
220    pub score: i64,
221
222    #[pyo3(get)]
223    pub reason: String,
224}
225#[pymethods]
226impl Score {
227    #[staticmethod]
228    pub fn response_type() -> ResponseType {
229        ResponseType::Score
230    }
231
232    #[staticmethod]
233    pub fn model_validate_json(json_string: String) -> Result<Score, TypeError> {
234        Ok(serde_json::from_str(&json_string)?)
235    }
236
237    #[staticmethod]
238    pub fn model_json_schema<'py>(py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
239        let schema = Score::get_structured_output_schema();
240        Ok(pythonize(py, &schema)?)
241    }
242
243    pub fn __str__(&self) -> String {
244        PyHelperFuncs::__str__(self)
245    }
246}
247
248impl StructuredOutput for Score {}
249
250#[pyclass(from_py_object)]
251#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
252pub enum ResponseType {
253    Score,
254    Pydantic,
255    #[default]
256    Null, // This is used when no response type is specified
257}
258
259impl Display for ResponseType {
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        match self {
262            ResponseType::Score => write!(f, "Score"),
263            ResponseType::Pydantic => write!(f, "Pydantic"),
264            ResponseType::Null => write!(f, "Null"),
265        }
266    }
267}
268
269// add conversion logic based on message conversion trait
270
271#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
272#[serde(untagged)]
273pub enum MessageNum {
274    OpenAIMessageV1(OpenAIChatMessage),
275    AnthropicMessageV1(AnthropicMessage),
276    GeminiContentV1(GeminiContent),
277
278    // this is a special case for Anthropic system messages
279    AnthropicSystemMessageV1(TextBlockParam),
280
281    // Raw JSON message - used internally in the agentic loop for provider-specific
282    // messages that don't fit the typed variants (e.g., OpenAI tool-call assistant messages).
283    // Must be LAST so serde tries typed variants first when deserializing.
284    RawV1(serde_json::Value),
285}
286
287impl MessageNum {
288    /// Checks if the message type matches the given provider
289    fn matches_provider(&self, provider: &Provider) -> bool {
290        matches!(
291            (self, provider),
292            (MessageNum::OpenAIMessageV1(_), Provider::OpenAI)
293                | (MessageNum::AnthropicMessageV1(_), Provider::Anthropic)
294                | (MessageNum::AnthropicSystemMessageV1(_), Provider::Anthropic)
295                | (MessageNum::GeminiContentV1(_), Provider::Google)
296                | (MessageNum::GeminiContentV1(_), Provider::Vertex)
297                | (MessageNum::GeminiContentV1(_), Provider::Gemini)
298        )
299    }
300
301    /// Converts the message to an openai message
302    /// This is only done for anthropic and gemini messages
303    /// openai message will return a failure if called on an openai message
304    /// Control flow should ensure this is only called on non-openai messages
305    fn to_openai_message(&self) -> Result<MessageNum, TypeError> {
306        match self {
307            MessageNum::AnthropicMessageV1(msg) => {
308                Ok(MessageNum::OpenAIMessageV1(msg.to_openai_message()?))
309            }
310            MessageNum::GeminiContentV1(msg) => {
311                Ok(MessageNum::OpenAIMessageV1(msg.to_openai_message()?))
312            }
313            _ => Err(TypeError::CantConvertSelf),
314        }
315    }
316
317    /// Converts to Anthropic message format
318    fn to_anthropic_message(&self) -> Result<MessageNum, TypeError> {
319        match self {
320            MessageNum::OpenAIMessageV1(msg) => {
321                Ok(MessageNum::AnthropicMessageV1(msg.to_anthropic_message()?))
322            }
323            MessageNum::GeminiContentV1(msg) => {
324                Ok(MessageNum::AnthropicMessageV1(msg.to_anthropic_message()?))
325            }
326            _ => Err(TypeError::CantConvertSelf),
327        }
328    }
329
330    /// Converts to Google Gemini message format
331    fn to_google_message(&self) -> Result<MessageNum, TypeError> {
332        match self {
333            MessageNum::OpenAIMessageV1(msg) => {
334                Ok(MessageNum::GeminiContentV1(msg.to_google_message()?))
335            }
336            MessageNum::AnthropicMessageV1(msg) => {
337                Ok(MessageNum::GeminiContentV1(msg.to_google_message()?))
338            }
339            _ => Err(TypeError::CantConvertSelf),
340        }
341    }
342
343    fn convert_message_to_provider_type(
344        &self,
345        provider: &Provider,
346    ) -> Result<MessageNum, TypeError> {
347        match provider {
348            Provider::OpenAI => self.to_openai_message(),
349            Provider::Anthropic => self.to_anthropic_message(),
350            Provider::Google => self.to_google_message(),
351            Provider::Vertex => self.to_google_message(),
352            Provider::Gemini => self.to_google_message(),
353            _ => Err(TypeError::UnsupportedProviderError),
354        }
355    }
356
357    pub fn convert_message(&mut self, provider: &Provider) -> Result<(), TypeError> {
358        // RawV1 is provider-specific raw JSON — skip conversion
359        if matches!(self, MessageNum::RawV1(_)) {
360            return Ok(());
361        }
362        // if message already matches provider, return Ok
363        if self.matches_provider(provider) {
364            return Ok(());
365        }
366        let converted = self.convert_message_to_provider_type(provider)?;
367        *self = converted;
368        Ok(())
369    }
370
371    pub fn anthropic_message_to_system_message(&mut self) -> Result<(), TypeError> {
372        match self {
373            MessageNum::AnthropicMessageV1(msg) => {
374                let text_param = msg.to_text_block_param()?;
375                *self = MessageNum::AnthropicSystemMessageV1(text_param);
376                Ok(())
377            }
378            _ => Err(TypeError::Error(
379                "Cannot convert non-AnthropicMessageV1 to system message".to_string(),
380            )),
381        }
382    }
383    pub fn role(&self) -> &str {
384        match self {
385            MessageNum::OpenAIMessageV1(msg) => &msg.role,
386            MessageNum::AnthropicMessageV1(msg) => &msg.role,
387            MessageNum::GeminiContentV1(msg) => &msg.role,
388            _ => "system",
389        }
390    }
391    pub fn bind(&self, name: &str, value: &str) -> Result<Self, TypeError> {
392        match self {
393            MessageNum::OpenAIMessageV1(msg) => {
394                let bound_msg = msg.bind(name, value)?;
395                Ok(MessageNum::OpenAIMessageV1(bound_msg))
396            }
397            MessageNum::AnthropicMessageV1(msg) => {
398                let bound_msg = msg.bind(name, value)?;
399                Ok(MessageNum::AnthropicMessageV1(bound_msg))
400            }
401            MessageNum::GeminiContentV1(msg) => {
402                let bound_msg = msg.bind(name, value)?;
403                Ok(MessageNum::GeminiContentV1(bound_msg))
404            }
405            _ => Ok(self.clone()),
406        }
407    }
408    pub fn bind_mut(&mut self, name: &str, value: &str) -> Result<(), TypeError> {
409        match self {
410            MessageNum::OpenAIMessageV1(msg) => msg.bind_mut(name, value),
411            MessageNum::AnthropicMessageV1(msg) => msg.bind_mut(name, value),
412            MessageNum::GeminiContentV1(msg) => msg.bind_mut(name, value),
413            _ => Ok(()),
414        }
415    }
416
417    pub(crate) fn extract_variables(&self) -> Vec<String> {
418        match self {
419            MessageNum::OpenAIMessageV1(msg) => msg.extract_variables(),
420            MessageNum::AnthropicMessageV1(msg) => msg.extract_variables(),
421            MessageNum::GeminiContentV1(msg) => msg.extract_variables(),
422            _ => vec![],
423        }
424    }
425
426    pub(crate) fn extract_media_variables(&self) -> Vec<String> {
427        match self {
428            MessageNum::OpenAIMessageV1(msg) => msg.extract_media_variables(),
429            MessageNum::AnthropicMessageV1(msg) => msg.extract_media_variables(),
430            MessageNum::GeminiContentV1(msg) => msg.extract_media_variables(),
431            MessageNum::AnthropicSystemMessageV1(t) => {
432                let mut out = Vec::new();
433                let regex = crate::traits::get_media_regex();
434                for cap in regex.captures_iter(&t.text) {
435                    if let Some(name) = cap.get(1) {
436                        out.push(name.as_str().to_string());
437                    }
438                }
439                out
440            }
441            _ => vec![],
442        }
443    }
444
445    pub fn split_media_placeholders(&mut self) -> Result<(), TypeError> {
446        match self {
447            MessageNum::OpenAIMessageV1(msg) => msg.split_media_placeholders(),
448            MessageNum::AnthropicMessageV1(msg) => msg.split_media_placeholders(),
449            MessageNum::GeminiContentV1(msg) => msg.split_media_placeholders(),
450            MessageNum::AnthropicSystemMessageV1(t) => {
451                let regex = crate::traits::get_media_regex();
452                if regex.is_match(&t.text) {
453                    Err(TypeError::MediaInSystemMessage)
454                } else {
455                    Ok(())
456                }
457            }
458            MessageNum::RawV1(_) => Ok(()),
459        }
460    }
461
462    pub fn bind_media_mut(
463        &mut self,
464        token: &str,
465        media: &crate::prompt::media::MediaRef,
466        provider: &crate::Provider,
467    ) -> Result<bool, TypeError> {
468        match self {
469            MessageNum::OpenAIMessageV1(msg) => msg.bind_media_mut(token, media, provider),
470            MessageNum::AnthropicMessageV1(msg) => msg.bind_media_mut(token, media, provider),
471            MessageNum::GeminiContentV1(msg) => msg.bind_media_mut(token, media, provider),
472            MessageNum::AnthropicSystemMessageV1(t) => {
473                if t.text.contains(token) {
474                    Err(TypeError::MediaInSystemMessage)
475                } else {
476                    Ok(false)
477                }
478            }
479            MessageNum::RawV1(_) => Ok(false),
480        }
481    }
482
483    pub fn to_bound_py_object<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
484        match self {
485            MessageNum::OpenAIMessageV1(msg) => {
486                let bound_msg = msg.clone().into_bound_py_any(py)?;
487                Ok(bound_msg)
488            }
489            MessageNum::AnthropicMessageV1(msg) => {
490                let bound_msg = msg.clone().into_bound_py_any(py)?;
491                Ok(bound_msg)
492            }
493            MessageNum::GeminiContentV1(msg) => {
494                let bound_msg = msg.clone().into_bound_py_any(py)?;
495                Ok(bound_msg)
496            }
497            MessageNum::AnthropicSystemMessageV1(msg) => {
498                // Convert to AnthropicMessage first
499                let anthropic_msg = AnthropicMessage::from_text(msg.text.clone(), self.role())?;
500                let bound_msg = anthropic_msg.into_bound_py_any(py)?;
501                Ok(bound_msg)
502            }
503            MessageNum::RawV1(v) => Ok(pythonize(py, v)?),
504        }
505    }
506
507    pub fn to_bound_openai_message<'py>(
508        &self,
509        py: Python<'py>,
510    ) -> Result<Bound<'py, OpenAIChatMessage>, TypeError> {
511        match self {
512            MessageNum::OpenAIMessageV1(msg) => {
513                let py_obj = Py::new(py, msg.clone())?;
514                let bound = py_obj.bind(py);
515                Ok(bound.clone())
516            }
517            _ => Err(TypeError::CantConvertSelf),
518        }
519    }
520
521    pub fn to_bound_gemini_message<'py>(
522        &self,
523        py: Python<'py>,
524    ) -> Result<Bound<'py, GeminiContent>, TypeError> {
525        match self {
526            MessageNum::GeminiContentV1(msg) => {
527                let py_obj = Py::new(py, msg.clone())?;
528                let bound = py_obj.bind(py);
529                Ok(bound.clone())
530            }
531            _ => Err(TypeError::CantConvertSelf),
532        }
533    }
534
535    pub fn to_bound_anthropic_message<'py>(
536        &self,
537        py: Python<'py>,
538    ) -> Result<Bound<'py, AnthropicMessage>, TypeError> {
539        match self {
540            MessageNum::AnthropicMessageV1(msg) => {
541                let py_obj = Py::new(py, msg.clone())?;
542                let bound = py_obj.bind(py);
543                Ok(bound.clone())
544            }
545            _ => Err(TypeError::CantConvertSelf),
546        }
547    }
548
549    pub fn is_system_message(&self) -> bool {
550        match self {
551            MessageNum::OpenAIMessageV1(msg) => {
552                msg.role == Role::Developer.to_string() || msg.role == Role::System.to_string()
553            }
554            MessageNum::AnthropicMessageV1(msg) => msg.role == Role::System.to_string(),
555            MessageNum::GeminiContentV1(msg) => msg.role == Role::Model.to_string(),
556            MessageNum::AnthropicSystemMessageV1(_) => true,
557            MessageNum::RawV1(v) => v
558                .get("role")
559                .and_then(|r| r.as_str())
560                .map(|r| r == "system" || r == "developer")
561                .unwrap_or(false),
562        }
563    }
564
565    pub fn is_user_message(&self) -> bool {
566        match self {
567            MessageNum::OpenAIMessageV1(msg) => msg.role == Role::User.to_string(),
568            MessageNum::AnthropicMessageV1(msg) => msg.role == Role::User.to_string(),
569            MessageNum::GeminiContentV1(msg) => msg.role == Role::User.to_string(),
570            MessageNum::AnthropicSystemMessageV1(_) => false,
571            MessageNum::RawV1(v) => v
572                .get("role")
573                .and_then(|r| r.as_str())
574                .map(|r| r == "user")
575                .unwrap_or(false),
576        }
577    }
578}
579
580#[derive(Debug, Clone)]
581pub enum ResponseContent {
582    OpenAI(Choice),
583    Google(Candidate),
584    Anthropic(ResponseContentBlock),
585    PredictResponse(PredictResponse),
586}
587
588#[pyclass(skip_from_py_object)]
589pub struct OpenAIMessageList {
590    pub messages: Vec<OpenAIChatMessage>,
591}
592
593#[pymethods]
594impl OpenAIMessageList {
595    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<OpenAIMessageIterator>> {
596        let iter = OpenAIMessageIterator {
597            inner: slf.messages.clone().into_iter(),
598        };
599        Py::new(slf.py(), iter)
600    }
601
602    pub fn __len__(&self) -> usize {
603        self.messages.len()
604    }
605
606    pub fn __getitem__(&self, index: isize) -> Result<OpenAIChatMessage, TypeError> {
607        let len = self.messages.len() as isize;
608        let normalized_index = if index < 0 { len + index } else { index };
609
610        if normalized_index < 0 || normalized_index >= len {
611            return Err(TypeError::Error(format!(
612                "Index {} out of range for list of length {}",
613                index, len
614            )));
615        }
616
617        Ok(self.messages[normalized_index as usize].clone())
618    }
619
620    pub fn __str__(&self) -> String {
621        PyHelperFuncs::__str__(&self.messages)
622    }
623
624    pub fn __repr__(&self) -> String {
625        self.__str__()
626    }
627}
628
629#[pyclass(skip_from_py_object)]
630pub struct OpenAIMessageIterator {
631    inner: std::vec::IntoIter<OpenAIChatMessage>,
632}
633
634#[pymethods]
635impl OpenAIMessageIterator {
636    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
637        slf
638    }
639
640    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<OpenAIChatMessage> {
641        slf.inner.next()
642    }
643}
644
645#[pyclass(skip_from_py_object)]
646pub struct AnthropicMessageList {
647    pub messages: Vec<AnthropicMessage>,
648}
649
650#[pymethods]
651impl AnthropicMessageList {
652    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<AnthropicMessageIterator>> {
653        let iter = AnthropicMessageIterator {
654            inner: slf.messages.clone().into_iter(),
655        };
656        Py::new(slf.py(), iter)
657    }
658
659    pub fn __len__(&self) -> usize {
660        self.messages.len()
661    }
662
663    pub fn __getitem__(&self, index: isize) -> Result<AnthropicMessage, TypeError> {
664        let len = self.messages.len() as isize;
665        let normalized_index = if index < 0 { len + index } else { index };
666
667        if normalized_index < 0 || normalized_index >= len {
668            return Err(TypeError::Error(format!(
669                "Index {} out of range for list of length {}",
670                index, len
671            )));
672        }
673
674        Ok(self.messages[normalized_index as usize].clone())
675    }
676
677    pub fn __str__(&self) -> String {
678        PyHelperFuncs::__str__(&self.messages)
679    }
680
681    pub fn __repr__(&self) -> String {
682        self.__str__()
683    }
684}
685
686#[pyclass(skip_from_py_object)]
687pub struct AnthropicMessageIterator {
688    inner: std::vec::IntoIter<AnthropicMessage>,
689}
690
691#[pymethods]
692impl AnthropicMessageIterator {
693    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
694        slf
695    }
696
697    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<AnthropicMessage> {
698        slf.inner.next()
699    }
700}
701
702#[pyclass(skip_from_py_object)]
703pub struct GeminiContentList {
704    pub messages: Vec<GeminiContent>,
705}
706
707#[pymethods]
708impl GeminiContentList {
709    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<GeminiContentIterator>> {
710        let iter = GeminiContentIterator {
711            inner: slf.messages.clone().into_iter(),
712        };
713        Py::new(slf.py(), iter)
714    }
715
716    pub fn __len__(&self) -> usize {
717        self.messages.len()
718    }
719
720    pub fn __getitem__(&self, index: isize) -> Result<GeminiContent, TypeError> {
721        let len = self.messages.len() as isize;
722        let normalized_index = if index < 0 { len + index } else { index };
723
724        if normalized_index < 0 || normalized_index >= len {
725            return Err(TypeError::Error(format!(
726                "Index {} out of range for list of length {}",
727                index, len
728            )));
729        }
730
731        Ok(self.messages[normalized_index as usize].clone())
732    }
733
734    pub fn __str__(&self) -> String {
735        PyHelperFuncs::__str__(&self.messages)
736    }
737
738    pub fn __repr__(&self) -> String {
739        self.__str__()
740    }
741}
742
743#[pyclass(skip_from_py_object)]
744pub struct GeminiContentIterator {
745    inner: std::vec::IntoIter<GeminiContent>,
746}
747
748#[pymethods]
749impl GeminiContentIterator {
750    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
751        slf
752    }
753
754    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<GeminiContent> {
755        slf.inner.next()
756    }
757}
potato_type/prompt/types.rs

potato_type/prompt/
types.rs
