potato_type/
traits.rs

1use crate::tools::ToolCallInfo;
2use crate::{
3    error::TypeError,
4    prompt::{MessageNum, ModelSettings, ResponseContent},
5    tools::AgentToolDefinition,
6    Provider,
7};
8use potato_util::utils::TokenLogProbs;
9use pyo3::prelude::*;
10use pyo3::types::PyList;
11use regex::Regex;
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14use std::sync::OnceLock;
15
16pub static VAR_REGEX: OnceLock<Regex> = OnceLock::new();
17pub fn get_var_regex() -> &'static Regex {
18    VAR_REGEX.get_or_init(|| Regex::new(r"\$?\{([a-zA-Z_][a-zA-Z0-9_]*)\}").unwrap())
19}
20
21pub static MEDIA_REGEX: OnceLock<Regex> = OnceLock::new();
22pub fn get_media_regex() -> &'static Regex {
23    MEDIA_REGEX.get_or_init(|| Regex::new(r"\$\{media:([a-zA-Z_][a-zA-Z0-9_]*)\}").unwrap())
24}
25
26pub(crate) fn split_text_on_media(text: &str, regex: &Regex) -> Vec<String> {
27    let mut out: Vec<String> = Vec::new();
28    let mut last = 0usize;
29    for m in regex.find_iter(text) {
30        if m.start() > last {
31            out.push(text[last..m.start()].to_string());
32        }
33        out.push(m.as_str().to_string());
34        last = m.end();
35    }
36    if last < text.len() {
37        out.push(text[last..].to_string());
38    }
39    out
40}
41
42pub(crate) fn extract_token_name(token: &str) -> String {
43    token
44        .strip_prefix("${media:")
45        .and_then(|s| s.strip_suffix('}'))
46        .unwrap_or("")
47        .to_string()
48}
49use crate::prompt::builder::ProviderRequest;
50/// Core trait that all message types must implement
51pub trait PromptMessageExt:
52    Send + Sync + Clone + Serialize + for<'de> Deserialize<'de> + PartialEq
53{
54    /// Bind a variable in the message content, returning a new instance
55    fn bind(&self, name: &str, value: &str) -> Result<Self, TypeError>
56    where
57        Self: Sized;
58
59    /// Bind a variable in-place
60    fn bind_mut(&mut self, name: &str, value: &str) -> Result<(), TypeError>;
61
62    /// Extract variables from the message content
63    fn extract_variables(&self) -> Vec<String>;
64
65    fn from_text(content: String, role: &str) -> Result<Self, TypeError>;
66
67    fn extract_media_variables(&self) -> Vec<String> {
68        Vec::new()
69    }
70
71    fn split_media_placeholders(&mut self) -> Result<(), TypeError> {
72        Ok(())
73    }
74
75    fn bind_media_mut(
76        &mut self,
77        _token: &str,
78        _media: &crate::prompt::media::MediaRef,
79        _provider: &crate::Provider,
80    ) -> Result<bool, TypeError> {
81        Ok(false)
82    }
83}
84
85/// Core trait that must be implemented for all request types
86pub trait RequestAdapter {
87    /// Returns all messages in the request
88    fn messages(&self) -> &[MessageNum];
89    /// Returns a mutable reference to the messages in the request
90    fn messages_mut(&mut self) -> &mut Vec<MessageNum>;
91    /// Returns all system instructions in the request
92    fn system_instructions(&self) -> Vec<&MessageNum>;
93    /// Returns the response JSON schema if set
94    fn response_json_schema(&self) -> Option<&Value>;
95    /// Inserts a message at the specified index (or at the start if None)
96    fn insert_message(&mut self, message: MessageNum, idx: Option<usize>) {
97        self.messages_mut().insert(idx.unwrap_or(0), message);
98    }
99    /// Prepends system instructions to the messages
100    fn preprend_system_instructions(&mut self, messages: Vec<MessageNum>) -> Result<(), TypeError>;
101
102    /// Returns the system instructions as a Python list
103    /// # Arguments
104    /// * `py` - The Python GIL token
105    /// # Returns
106    /// Returns a Python list of system instruction messages
107    fn get_py_system_instructions<'py>(
108        &self,
109        py: Python<'py>,
110    ) -> Result<Bound<'py, PyList>, TypeError>;
111    /// Returns the model settings for the request (python object)
112    fn model_settings<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError>;
113
114    /// Converts the request to a JSON value for sending to the provider
115    fn to_request_body(&self) -> Result<Value, TypeError>;
116    /// Checks if the request matches the given provider
117    fn match_provider(&self, provider: &Provider) -> bool;
118    /// Builds a provider-specific request enum from the given parameters
119    /// The ProviderRequest enum encapsulates all supported provider request types and is an
120    /// attribute of the Prompt struct. ProviderRequest is built on instantiation of the Prompt
121    fn build_provider_enum(
122        messages: Vec<MessageNum>,
123        system_instructions: Vec<MessageNum>,
124        model: String,
125        settings: ModelSettings,
126        response_json_schema: Option<Value>,
127    ) -> Result<ProviderRequest, TypeError>;
128
129    /// Sets the response JSON schema for the request
130    /// Typically used as part of workflows when adding tasks
131    fn set_response_json_schema(&mut self, response_json_schema: Option<Value>) -> ();
132
133    /// Adds tools to the request
134    fn add_tools(&mut self, tools: Vec<AgentToolDefinition>) -> Result<(), TypeError>;
135}
136
137pub trait ResponseAdapter {
138    /// Returns a string representation of the response
139    fn __str__(&self) -> String;
140
141    /// Checks if the response is empty
142    fn is_empty(&self) -> bool;
143
144    /// Converts the response to a Python object
145    fn to_bound_py_object<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError>;
146
147    /// Returns the response ID
148    fn id(&self) -> &str;
149
150    /// Converts the response to a vector of MessageNum
151    fn to_message_num(&self) -> Result<Vec<MessageNum>, TypeError>;
152
153    // Get the token usage as a Python object
154    fn usage<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError>;
155
156    /// Retrieves the first content choice from the response
157    fn get_content(&self) -> ResponseContent;
158
159    /// Retrieves the log probabilities from the response
160    fn get_log_probs(&self) -> Vec<TokenLogProbs>;
161
162    /// Returns the structured output of the response
163    /// For all response types the flow is as follows:
164    /// 1. Check if the response has content (string/text)
165    /// 2. If no content, return Python None
166    /// 3. If content exists, check if an output_type/model is provided
167    /// 4. If output_type/model is provided, attempt to convert the content to that type
168    /// 5. If conversion fails, attempt to construct a generic Python object from the content
169    /// 6. If no output_type/model is provided, return the content as a generic Python object
170    /// # Arguments
171    /// * `py`: The Python GIL token
172    /// * `output_type`: An optional Python type/model to convert the content into. This can be a pydantic model or any object
173    ///   that implements model_validate_json that can parse from a JSON string.
174    /// # Returns
175    /// * `Result<Bound<'py, PyAny>, TypeError>`: The structured output as a Python object or an error
176    fn structured_output<'py>(
177        &self,
178        py: Python<'py>,
179        output_type: Option<&Bound<'py, PyAny>>,
180    ) -> Result<Bound<'py, PyAny>, TypeError>;
181
182    /// Returns the structured output value as a serde_json::Value
183    fn structured_output_value(&self) -> Option<Value>;
184
185    /// Returns any tool calls made in the response, if applicable
186    fn tool_call_output(&self) -> Option<Value>;
187
188    /// Returns the output text of the response if available
189    fn response_text(&self) -> String;
190
191    /// Extracts tool calls from the response, if any.
192    /// Returns None if the response contains no tool calls.
193    fn extract_tool_calls(&self) -> Option<Vec<crate::tools::ToolCall>> {
194        None
195    }
196    fn model_name(&self) -> Option<&str>;
197
198    fn finish_reason(&self) -> Option<&str>;
199
200    fn input_tokens(&self) -> Option<i64>;
201
202    fn output_tokens(&self) -> Option<i64>;
203
204    fn total_tokens(&self) -> Option<i64>;
205
206    fn get_tool_calls(&self) -> Vec<ToolCallInfo>;
207}
208
209pub trait MessageResponseExt {
210    fn to_message_num(&self) -> Result<MessageNum, TypeError>;
211}
212
213pub trait MessageFactory: Sized {
214    fn from_text(content: String, role: &str) -> Result<Self, TypeError>;
215}
216
217/// Trait for converting between different provider message formats
218///
219/// This trait enables conversion of messages between different LLM provider formats
220/// (e.g., Anthropic MessageParam ↔ Google GeminiContent ↔ OpenAI ChatMessage).
221///
222/// Currently focused on text content conversion, with support for other content
223/// types planned for future implementation.
224pub trait MessageConversion {
225    /// Convert this message to an Anthropic MessageParam
226    ///
227    /// # Errors
228    /// Returns `TypeError::UnsupportedConversion` if the message contains
229    /// content types that cannot be represented in Anthropic's format
230    fn to_anthropic_message(
231        &self,
232    ) -> Result<crate::anthropic::v1::request::MessageParam, TypeError>;
233
234    /// Convert this message to a Google GeminiContent
235    ///
236    /// # Errors
237    /// Returns `TypeError::UnsupportedConversion` if the message contains
238    /// content types that cannot be represented in Google's format
239    fn to_google_message(
240        &self,
241    ) -> Result<crate::google::v1::generate::request::GeminiContent, TypeError>;
242
243    /// Convert this message to an OpenAI ChatMessage
244    ///
245    /// # Errors
246    /// Returns `TypeError::UnsupportedConversion` if the message contains
247    /// content types that cannot be represented in OpenAI's format
248    fn to_openai_message(&self)
249        -> Result<crate::openai::v1::chat::request::ChatMessage, TypeError>;
250}
potato_type/traits.rs

potato_type/
traits.rs
