potato_util/
utils.rs

1use crate::error::UtilError;
2
3use colored_json::{Color, ColorMode, ColoredFormatter, PrettyFormatter, Styler};
4use pyo3::prelude::*;
5
6use pyo3::types::{PyAny, PyDict, PyList};
7use pyo3::IntoPyObjectExt;
8use serde::Serialize;
9use serde_json::Value;
10use serde_json::Value::{Null, Object};
11use std::ops::RangeInclusive;
12use std::path::Path;
13use uuid::Uuid;
14pub fn create_uuid7() -> String {
15    Uuid::now_v7().to_string()
16}
17use pythonize::{depythonize, pythonize};
18use tracing::warn;
19pub struct PyHelperFuncs {}
20
21impl PyHelperFuncs {
22    /// Convert any type implementing `IntoPyObject` to a Python object
23    /// # Arguments
24    /// * `py` - A Python interpreter instance
25    /// * `object` - A reference to an object implementing `IntoPyObject`
26    /// # Returns
27    /// * `Result<Bound<'py, PyAny>, UtilError>` - A result containing the Python object or an error
28    pub fn to_bound_py_object<'py, T>(
29        py: Python<'py>,
30        object: &T,
31    ) -> Result<Bound<'py, PyAny>, UtilError>
32    where
33        T: IntoPyObject<'py> + Clone,
34    {
35        Ok(object.clone().into_bound_py_any(py)?)
36    }
37    pub fn __str__<T: Serialize>(object: T) -> String {
38        match ColoredFormatter::with_styler(
39            PrettyFormatter::default(),
40            Styler {
41                key: Color::Rgb(75, 57, 120).foreground(),
42                string_value: Color::Rgb(4, 205, 155).foreground(),
43                float_value: Color::Rgb(4, 205, 155).foreground(),
44                integer_value: Color::Rgb(4, 205, 155).foreground(),
45                bool_value: Color::Rgb(4, 205, 155).foreground(),
46                nil_value: Color::Rgb(4, 205, 155).foreground(),
47                ..Default::default()
48            },
49        )
50        .to_colored_json(&object, ColorMode::On)
51        {
52            Ok(json) => json,
53            Err(e) => format!("Failed to serialize to json: {e}"),
54        }
55        // serialize the struct to a string
56    }
57
58    pub fn __json__<T: Serialize>(object: T) -> String {
59        match serde_json::to_string_pretty(&object) {
60            Ok(json) => json,
61            Err(e) => format!("Failed to serialize to json: {e}"),
62        }
63    }
64
65    /// Save a struct to a JSON file
66    ///
67    /// # Arguments
68    ///
69    /// * `model` - A reference to a struct that implements the `Serialize` trait
70    /// * `path` - A reference to a `Path` object that holds the path to the file
71    ///
72    /// # Returns
73    ///
74    /// A `Result` containing `()` or a `UtilError`
75    ///
76    /// # Errors
77    ///
78    /// This function will return an error if:
79    /// - The struct cannot be serialized to a string
80    pub fn save_to_json<T>(model: T, path: &Path) -> Result<(), UtilError>
81    where
82        T: Serialize,
83    {
84        // serialize the struct to a string
85        let json =
86            serde_json::to_string_pretty(&model).map_err(|_| UtilError::SerializationError)?;
87
88        // ensure .json extension
89        let path = path.with_extension("json");
90
91        if !path.exists() {
92            // ensure path exists, create if not
93            let parent_path = path.parent().ok_or(UtilError::GetParentPathError)?;
94            if !parent_path.as_os_str().is_empty() {
95                std::fs::create_dir_all(parent_path)
96                    .map_err(|_| UtilError::CreateDirectoryError)?;
97            }
98        }
99
100        std::fs::write(path, json).map_err(|_| UtilError::WriteError)?;
101
102        Ok(())
103    }
104}
105
106pub fn vec_to_py_object<'py>(
107    py: Python<'py>,
108    vec: &Vec<Value>,
109) -> Result<Bound<'py, PyList>, UtilError> {
110    let py_list = PyList::empty(py);
111    for item in vec {
112        let py_item = pythonize(py, item)?;
113        py_list.append(py_item)?;
114    }
115    Ok(py_list)
116}
117
118pub fn version() -> String {
119    env!("CARGO_PKG_VERSION").to_string()
120}
121
122pub fn update_serde_value(value: &mut Value, key: &str, new_value: Value) -> Result<(), UtilError> {
123    if let Value::Object(map) = value {
124        map.insert(key.to_string(), new_value);
125        Ok(())
126    } else {
127        Err(UtilError::RootMustBeObjectError)
128    }
129}
130
131/// Updates a serde_json::Value object with another serde_json::Value object.
132/// Both types must be of the `Object` variant.
133/// If a key in the source object does not exist in the destination object,
134/// it will be added with the value from the source object.
135/// # Arguments
136/// * `dest` - A mutable reference to the destination serde_json::Value object.
137/// * `src` - A reference to the source serde_json::Value object.
138/// # Returns
139/// * `Ok(())` if the update was successful.
140/// * `Err(UtilError::RootMustBeObjectError)` if either `dest` or `src` is not an `Object`.
141pub fn update_serde_map_with(
142    dest: &mut serde_json::Value,
143    src: &serde_json::Value,
144) -> Result<(), UtilError> {
145    match (dest, src) {
146        (&mut Object(ref mut map_dest), Object(ref map_src)) => {
147            // map_dest and map_src both are Map<String, Value>
148            for (key, value) in map_src {
149                // if key is not in map_dest, create a Null object
150                // then only, update the value
151                *map_dest.entry(key.clone()).or_insert(Null) = value.clone();
152            }
153            Ok(())
154        }
155        (_, _) => Err(UtilError::RootMustBeObjectError),
156    }
157}
158
159/// Extracts a string value from a Python object.
160pub fn extract_string_value(py_value: &Bound<'_, PyAny>) -> Result<String, UtilError> {
161    // Try to extract as string first (most common case)
162    if let Ok(string_val) = py_value.extract::<String>() {
163        return Ok(string_val);
164    }
165    // Try to extract as boolean
166    if let Ok(bool_val) = py_value.extract::<bool>() {
167        return Ok(bool_val.to_string());
168    }
169
170    // Try to extract as integer
171    if let Ok(int_val) = py_value.extract::<i64>() {
172        return Ok(int_val.to_string());
173    }
174
175    // Try to extract as float
176    if let Ok(float_val) = py_value.extract::<f64>() {
177        return Ok(float_val.to_string());
178    }
179
180    // For complex objects, convert to JSON but extract the value without quotes
181    let json_value = depythonize(py_value)?;
182
183    match json_value {
184        Value::String(s) => Ok(s),
185        Value::Number(n) => Ok(n.to_string()),
186        Value::Bool(b) => Ok(b.to_string()),
187        Value::Null => Ok("null".to_string()),
188        _ => {
189            // For arrays and objects, serialize to JSON string
190            let json_string = serde_json::to_string(&json_value)?;
191            Ok(json_string)
192        }
193    }
194}
195
196#[pyclass(from_py_object)]
197#[derive(Debug, Serialize, Clone)]
198pub struct TokenLogProbs {
199    #[pyo3(get)]
200    pub token: String,
201
202    #[pyo3(get)]
203    pub logprob: f64,
204}
205
206#[pyclass(from_py_object)]
207#[derive(Debug, Serialize, Clone)]
208pub struct ResponseLogProbs {
209    #[pyo3(get)]
210    pub tokens: Vec<TokenLogProbs>,
211}
212
213#[pymethods]
214impl ResponseLogProbs {
215    pub fn __str__(&self) -> String {
216        PyHelperFuncs::__str__(self)
217    }
218}
219
220/// Calculate a weighted score base on the log probabilities of tokens 1-5.
221pub fn calculate_weighted_score(log_probs: &[TokenLogProbs]) -> Result<Option<f64>, UtilError> {
222    let score_range = RangeInclusive::new(1, 5);
223    let mut score_probs = Vec::new();
224    let mut weighted_sum = 0.0;
225    let mut total_prob = 0.0;
226
227    for log_prob in log_probs {
228        let token = log_prob.token.parse::<u64>().ok();
229
230        if let Some(token_val) = token {
231            if score_range.contains(&token_val) {
232                let prob = log_prob.logprob.exp();
233                score_probs.push((token_val, prob));
234            }
235        }
236    }
237
238    for (score, logprob) in score_probs {
239        weighted_sum += score as f64 * logprob;
240        total_prob += logprob;
241    }
242
243    if total_prob > 0.0 {
244        Ok(Some(weighted_sum / total_prob))
245    } else {
246        Ok(None)
247    }
248}
249
250/// Generic function to convert text to a structured output model
251/// It is expected that output_model is a pydantic model or a potatohead type that implements serde json deserialization
252/// via model_validate_json method.
253/// Flow:
254/// 1. Attempt to validate the model using model_validate_json
255/// 2. If validation fails, attempt to parse the text as JSON and convert to python object
256/// # Arguments
257/// * `py` - A Python interpreter instance
258/// * `text` - The text to be converted (typically from an LLM response that returns structured output)
259/// * `output_model` - A bound python object representing the output model
260/// # Returns
261/// * `Result<Bound<'py, PyAny>, UtilError>` - A result containing the structured output or an error
262pub fn convert_text_to_structured_output<'py>(
263    py: Python<'py>,
264    text: String,
265    output_model: &Bound<'py, PyAny>,
266) -> Result<Bound<'py, PyAny>, UtilError> {
267    let output = output_model.call_method1("model_validate_json", (&text,));
268    match output {
269        Ok(obj) => {
270            // Successfully validated the model
271            Ok(obj)
272        }
273        Err(err) => {
274            // Model validation failed
275            // convert string to json and then to python object
276            warn!(
277                "Failed to validate model: {}, Attempting fallback to JSON parsing",
278                err
279            );
280            let val = serde_json::from_str::<serde_json::Value>(&text)?;
281            Ok(pythonize(py, &val)?)
282        }
283    }
284}
285
286pub fn is_pydantic_basemodel(py: Python, obj: &Bound<'_, PyAny>) -> Result<bool, UtilError> {
287    let pydantic = match py.import("pydantic") {
288        Ok(module) => module,
289        // return false if pydantic cannot be imported
290        Err(_) => return Ok(false),
291    };
292
293    let basemodel = pydantic.getattr("BaseModel")?;
294
295    // check if context is a pydantic model
296    let is_basemodel = obj
297        .is_instance(&basemodel)
298        .map_err(|e| UtilError::FailedToCheckPydanticModel(e.to_string()))?;
299
300    Ok(is_basemodel)
301}
302
303fn process_dict_with_nested_models(
304    py: Python<'_>,
305    dict: &Bound<'_, PyAny>,
306) -> Result<Value, UtilError> {
307    let py_dict = dict.cast::<PyDict>()?;
308    let mut result = serde_json::Map::new();
309
310    for (key, value) in py_dict.iter() {
311        let key_str: String = key.extract()?;
312        let processed_value = depythonize_object_to_value(py, &value)?;
313        result.insert(key_str, processed_value);
314    }
315
316    Ok(Value::Object(result))
317}
318
319pub fn depythonize_object_to_value<'py>(
320    py: Python<'py>,
321    value: &Bound<'py, PyAny>,
322) -> Result<Value, UtilError> {
323    let py_value = if is_pydantic_basemodel(py, value)? {
324        let model = value.call_method0("model_dump")?;
325        depythonize(&model)?
326    } else if value.is_instance_of::<PyDict>() {
327        process_dict_with_nested_models(py, value)?
328    } else {
329        depythonize(value)?
330    };
331    Ok(py_value)
332}
333
334/// Helper function to extract result from LLM response text
335/// If an output model is provided, it will attempt to convert the text to the structured output
336/// using the provided model. If no model is provided, it will attempt to convert the response to an appropriate
337/// Python type directly.
338/// # Arguments
339/// * `py` - A Python interpreter instance
340/// * `text` - The text to be converted (typically from an LLM response)
341/// * `output_model` - An optional bound python object representing the output model
342/// # Returns
343/// * `Result<Bound<'py, PyAny>, UtilError>` - A result containing the structured output or an error
344pub fn construct_structured_response<'py>(
345    py: Python<'py>,
346    text: String,
347    output_model: Option<&Bound<'py, PyAny>>,
348) -> Result<Bound<'py, PyAny>, UtilError> {
349    match output_model {
350        Some(model) => convert_text_to_structured_output(py, text, model),
351        None => {
352            // No output model provided, return the text as a Python string
353            let val = Value::String(text);
354            Ok(pythonize(py, &val)?)
355        }
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362    #[test]
363    fn test_calculate_weighted_score() {
364        let log_probs = vec![
365            TokenLogProbs {
366                token: "1".into(),
367                logprob: 0.9,
368            },
369            TokenLogProbs {
370                token: "2".into(),
371                logprob: 0.8,
372            },
373            TokenLogProbs {
374                token: "3".into(),
375                logprob: 0.7,
376            },
377        ];
378
379        let result = calculate_weighted_score(&log_probs);
380        assert!(result.is_ok());
381
382        let val = result.unwrap().unwrap();
383        // round to int
384        assert_eq!(val.round(), 2.0);
385    }
386    #[test]
387    fn test_calculate_weighted_score_empty() {
388        let log_probs: Vec<TokenLogProbs> = vec![];
389        let result = calculate_weighted_score(&log_probs);
390        assert!(result.is_ok());
391        assert_eq!(result.unwrap(), None);
392    }
393}
potato_util/utils.rs

potato_util/
utils.rs
