scouter_types/queue/
types.rs

1use crate::error::TypeError;
2use crate::json_to_pyobject_value;
3use crate::util::{is_pydantic_model, pyobject_to_json};
4use crate::PyHelperFuncs;
5use chrono::DateTime;
6use chrono::Utc;
7use potato_head::create_uuid7;
8use potato_head::Prompt;
9use pyo3::prelude::*;
10use pyo3::types::{PyDict, PyFloat, PyInt, PyList, PyString};
11use pyo3::IntoPyObjectExt;
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14use std::collections::HashMap;
15use std::fmt;
16use std::fmt::Display;
17use std::fmt::Formatter;
18
19#[pyclass]
20#[derive(Clone, Debug, Serialize)]
21pub enum EntityType {
22    Feature,
23    Metric,
24    LLM,
25}
26
27#[pyclass]
28#[derive(Clone, Serialize, Deserialize, Debug)]
29pub struct IntFeature {
30    pub name: String,
31    pub value: i64,
32}
33
34#[pymethods]
35impl IntFeature {
36    pub fn __str__(&self) -> String {
37        PyHelperFuncs::__str__(self)
38    }
39}
40
41impl IntFeature {
42    pub fn to_float(&self) -> f64 {
43        self.value as f64
44    }
45}
46
47#[pyclass]
48#[derive(Clone, Serialize, Deserialize, Debug)]
49pub struct FloatFeature {
50    pub name: String,
51    pub value: f64,
52}
53
54#[pymethods]
55impl FloatFeature {
56    pub fn __str__(&self) -> String {
57        PyHelperFuncs::__str__(self)
58    }
59}
60
61#[pyclass]
62#[derive(Clone, Serialize, Deserialize, Debug)]
63pub struct StringFeature {
64    pub name: String,
65    pub value: String,
66}
67
68#[pymethods]
69impl StringFeature {
70    pub fn __str__(&self) -> String {
71        PyHelperFuncs::__str__(self)
72    }
73}
74
75impl StringFeature {
76    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
77        feature_map
78            .features
79            .get(&self.name)
80            .and_then(|feat_map| {
81                feat_map
82                    .get(&self.value)
83                    .or_else(|| feat_map.get("missing"))
84            })
85            .map(|&val| val as f64)
86            .ok_or(TypeError::MissingStringValueError)
87    }
88}
89
90#[pyclass]
91#[derive(Clone, Serialize, Deserialize, Debug)]
92pub enum Feature {
93    Int(IntFeature),
94    Float(FloatFeature),
95    String(StringFeature),
96}
97
98#[pymethods]
99impl Feature {
100    #[new]
101    /// Parses a value to it's corresponding feature type.
102    /// PyFLoat -> FloatFeature
103    /// PyInt -> IntFeature
104    /// PyString -> StringFeature
105    /// # Arguments
106    /// * `name` - The name of the feature.
107    /// * `feature` - The value of the feature, which can be a PyFloat
108    /// # Returns
109    /// * `Feature` - The corresponding feature type.
110    /// # Errors
111    /// * `TypeError` - If the feature type is not supported.
112    pub fn new(name: &str, feature: Bound<'_, PyAny>) -> Result<Self, TypeError> {
113        // check python type
114        if feature.is_instance_of::<PyFloat>() {
115            let value: f64 = feature.extract().unwrap();
116            Ok(Feature::Float(FloatFeature {
117                name: name.into(),
118                value,
119            }))
120        } else if feature.is_instance_of::<PyInt>() {
121            let value: i64 = feature.extract().unwrap();
122            Ok(Feature::Int(IntFeature {
123                name: name.into(),
124                value,
125            }))
126        } else if feature.is_instance_of::<PyString>() {
127            let value: String = feature.extract().unwrap();
128            Ok(Feature::String(StringFeature {
129                name: name.into(),
130                value,
131            }))
132        } else {
133            Err(TypeError::UnsupportedFeatureTypeError(
134                feature.get_type().name()?.to_string(),
135            ))
136        }
137    }
138
139    #[staticmethod]
140    pub fn int(name: String, value: i64) -> Self {
141        Feature::Int(IntFeature { name, value })
142    }
143
144    #[staticmethod]
145    pub fn float(name: String, value: f64) -> Self {
146        Feature::Float(FloatFeature { name, value })
147    }
148
149    #[staticmethod]
150    pub fn string(name: String, value: String) -> Self {
151        Feature::String(StringFeature { name, value })
152    }
153
154    #[staticmethod]
155    pub fn categorical(name: String, value: String) -> Self {
156        Feature::String(StringFeature { name, value })
157    }
158
159    pub fn __str__(&self) -> String {
160        PyHelperFuncs::__str__(self)
161    }
162}
163
164impl Feature {
165    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
166        match self {
167            Feature::Int(feature) => Ok(feature.to_float()),
168            Feature::Float(feature) => Ok(feature.value),
169            Feature::String(feature) => feature.to_float(feature_map),
170        }
171    }
172
173    pub fn name(&self) -> &str {
174        match self {
175            Feature::Int(feature) => &feature.name,
176            Feature::Float(feature) => &feature.name,
177            Feature::String(feature) => &feature.name,
178        }
179    }
180
181    pub fn to_usize(&self, feature_map: &FeatureMap) -> Result<usize, TypeError> {
182        match self {
183            Feature::Int(f) => Ok(f.value as usize),
184            Feature::Float(f) => Ok(f.value as usize),
185            Feature::String(f) => Ok(f.to_float(feature_map)? as usize),
186        }
187    }
188}
189
190impl Display for Feature {
191    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
192        match self {
193            Feature::Int(feature) => write!(f, "{}", feature.value),
194            Feature::Float(feature) => write!(f, "{}", feature.value),
195            Feature::String(feature) => write!(f, "{}", feature.value),
196        }
197    }
198}
199
200#[pyclass]
201#[derive(Clone, Debug, Serialize)]
202pub struct Features {
203    #[pyo3(get)]
204    pub features: Vec<Feature>,
205
206    #[pyo3(get)]
207    pub entity_type: EntityType,
208}
209
210#[pymethods]
211impl Features {
212    #[new]
213    /// Creates a new Features instance.
214    /// A user may supply either a list of features or a single feature.
215    /// Extract features into a Vec<Feature>
216    /// Extraction follows the following rules:
217    /// 1. Check if Pylist, if so, extract to Vec<Feature>
218    /// 2. Check if PyDict, if so, iterate over each key-value pair and create a Feature
219    /// 3. If neither, return an error
220    /// # Arguments
221    /// * `features` - A Python object that can be a list of Feature instances or
222    ///               a dictionary of key-value pairs where keys are feature names
223    /// # Returns
224    /// * `Features` - A new Features instance containing the extracted features.
225    pub fn new(features: Bound<'_, PyAny>) -> Result<Self, TypeError> {
226        let features = if features.is_instance_of::<PyList>() {
227            features
228                .downcast::<PyList>()
229                .unwrap()
230                .iter()
231                .map(|item| item.extract::<Feature>().unwrap())
232                .collect()
233        } else if features.is_instance_of::<PyDict>() {
234            features
235                .downcast::<PyDict>()
236                .unwrap()
237                .iter()
238                .map(|(key, value)| {
239                    Feature::new(&key.extract::<String>().unwrap(), value.clone()).unwrap()
240                })
241                .collect()
242        } else {
243            Err(TypeError::UnsupportedFeaturesTypeError(
244                features.get_type().name()?.to_string(),
245            ))?
246        };
247        Ok(Features {
248            features,
249            entity_type: EntityType::Feature,
250        })
251    }
252
253    pub fn __str__(&self) -> String {
254        PyHelperFuncs::__str__(self)
255    }
256}
257
258impl Features {
259    pub fn iter(&self) -> std::slice::Iter<'_, Feature> {
260        self.features.iter()
261    }
262}
263
264#[pyclass]
265#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
266pub struct FeatureMap {
267    #[pyo3(get)]
268    pub features: HashMap<String, HashMap<String, usize>>,
269}
270
271#[pymethods]
272impl FeatureMap {
273    pub fn __str__(&self) -> String {
274        // serialize the struct to a string
275        PyHelperFuncs::__str__(self)
276    }
277}
278
279#[pyclass]
280#[derive(Clone, Serialize, Debug)]
281pub struct Metric {
282    pub name: String,
283    pub value: f64,
284}
285
286#[pymethods]
287impl Metric {
288    #[new]
289    pub fn new(name: String, value: Bound<'_, PyAny>) -> Self {
290        let value = if value.is_instance_of::<PyFloat>() {
291            value.extract::<f64>().unwrap()
292        } else if value.is_instance_of::<PyInt>() {
293            value.extract::<i64>().unwrap() as f64
294        } else {
295            panic!(
296                "Unsupported metric type: {}",
297                value.get_type().name().unwrap()
298            );
299        };
300        let lowercase_name = name.to_lowercase();
301        Metric {
302            name: lowercase_name,
303            value,
304        }
305    }
306
307    pub fn __str__(&self) -> String {
308        PyHelperFuncs::__str__(self)
309    }
310}
311
312impl Metric {
313    pub fn new_rs(name: String, value: f64) -> Self {
314        Metric { name, value }
315    }
316}
317
318#[pyclass]
319#[derive(Clone, Serialize, Debug)]
320pub struct Metrics {
321    #[pyo3(get)]
322    pub metrics: Vec<Metric>,
323
324    #[pyo3(get)]
325    pub entity_type: EntityType,
326}
327
328#[pymethods]
329impl Metrics {
330    #[new]
331    pub fn new(metrics: Bound<'_, PyAny>) -> Result<Self, TypeError> {
332        let metrics = if metrics.is_instance_of::<PyList>() {
333            metrics
334                .downcast::<PyList>()
335                .unwrap()
336                .iter()
337                .map(|item| item.extract::<Metric>().unwrap())
338                .collect()
339        } else if metrics.is_instance_of::<PyDict>() {
340            metrics
341                .downcast::<PyDict>()
342                .unwrap()
343                .iter()
344                .map(|(key, value)| Metric::new(key.extract().unwrap(), value))
345                .collect()
346        } else {
347            Err(TypeError::UnsupportedMetricsTypeError(
348                metrics.get_type().name()?.to_string(),
349            ))?
350        };
351        Ok(Metrics {
352            metrics,
353            entity_type: EntityType::Metric,
354        })
355    }
356    pub fn __str__(&self) -> String {
357        PyHelperFuncs::__str__(self)
358    }
359}
360
361impl Metrics {
362    pub fn iter(&self) -> std::slice::Iter<'_, Metric> {
363        self.metrics.iter()
364    }
365}
366
367#[pyclass]
368#[derive(Clone, Serialize, Debug)]
369pub struct LLMRecord {
370    pub uid: String,
371
372    pub space: String,
373
374    pub name: String,
375
376    pub version: String,
377
378    pub created_at: DateTime<Utc>,
379
380    pub context: Value,
381
382    pub score: Value,
383
384    pub prompt: Option<Value>,
385
386    #[pyo3(get)]
387    pub entity_type: EntityType,
388}
389
390#[pymethods]
391impl LLMRecord {
392    #[new]
393    #[pyo3(signature = (
394        context,
395        prompt=None,
396    ))]
397
398    /// Creates a new LLMRecord instance.
399    /// The context is either a python dictionary or a pydantic basemodel.
400    pub fn new(
401        py: Python<'_>,
402        context: Bound<'_, PyAny>,
403        prompt: Option<Bound<'_, PyAny>>,
404    ) -> Result<Self, TypeError> {
405        // check if context is a PyDict or PyObject(Pydantic model)
406        let context_val = if context.is_instance_of::<PyDict>() {
407            pyobject_to_json(&context)?
408        } else if is_pydantic_model(py, &context)? {
409            // Dump pydantic model to dictionary
410            let model = context.call_method0("model_dump")?;
411
412            // Serialize the dictionary to JSON
413            pyobject_to_json(&model)?
414        } else {
415            Err(TypeError::MustBeDictOrBaseModel)?
416        };
417
418        let prompt: Option<Value> = match prompt {
419            Some(p) => {
420                if p.is_instance_of::<Prompt>() {
421                    let prompt = p.extract::<Prompt>()?;
422                    Some(serde_json::to_value(prompt)?)
423                } else {
424                    Some(pyobject_to_json(&p)?)
425                }
426            }
427            None => None,
428        };
429
430        Ok(LLMRecord {
431            uid: create_uuid7(),
432            created_at: Utc::now(),
433            space: String::new(),
434            name: String::new(),
435            version: String::new(),
436            context: context_val,
437            score: Value::Null,
438            prompt,
439            entity_type: EntityType::LLM,
440        })
441    }
442
443    #[getter]
444    pub fn context<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
445        Ok(json_to_pyobject_value(py, &self.context)?
446            .into_bound_py_any(py)?
447            .clone())
448    }
449}
450
451impl LLMRecord {
452    pub fn new_rs(context: Option<Value>, prompt: Option<Value>) -> Self {
453        LLMRecord {
454            context: context.unwrap_or(Value::Object(serde_json::Map::new())),
455            prompt,
456            entity_type: EntityType::LLM,
457            uid: create_uuid7(),
458            created_at: Utc::now(),
459            space: String::new(),
460            name: String::new(),
461            version: String::new(),
462            score: Value::Null,
463        }
464    }
465
466    pub fn __str__(&self) -> String {
467        PyHelperFuncs::__str__(self)
468    }
469}
470
471#[derive(Debug)]
472pub enum QueueItem {
473    Features(Features),
474    Metrics(Metrics),
475    LLM(Box<LLMRecord>),
476}
477
478impl QueueItem {
479    /// Helper for extracting an Entity from a Python object
480    pub fn from_py_entity(entity: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
481        let entity_type = entity.getattr("entity_type")?.extract::<EntityType>()?;
482
483        match entity_type {
484            EntityType::Feature => {
485                let features = entity.extract::<Features>()?;
486                Ok(QueueItem::Features(features))
487            }
488            EntityType::Metric => {
489                let metrics = entity.extract::<Metrics>()?;
490                Ok(QueueItem::Metrics(metrics))
491            }
492            EntityType::LLM => {
493                // LLM is not supported in this context
494                let llm = entity.extract::<LLMRecord>()?;
495                Ok(QueueItem::LLM(Box::new(llm)))
496            }
497        }
498    }
499}
500
501pub trait QueueExt: Send + Sync {
502    fn metrics(&self) -> &Vec<Metric>;
503    fn features(&self) -> &Vec<Feature>;
504    fn llm_records(&self) -> Vec<&LLMRecord>;
505}
506
507impl QueueExt for Features {
508    fn metrics(&self) -> &Vec<Metric> {
509        // this is not a real implementation, just a placeholder
510        // to satisfy the trait bound
511        static EMPTY: Vec<Metric> = Vec::new();
512        &EMPTY
513    }
514
515    fn features(&self) -> &Vec<Feature> {
516        &self.features
517    }
518
519    fn llm_records(&self) -> Vec<&LLMRecord> {
520        // this is not a real implementation, just a placeholder
521        // to satisfy the trait bound
522        vec![]
523    }
524}
525
526impl QueueExt for Metrics {
527    fn metrics(&self) -> &Vec<Metric> {
528        &self.metrics
529    }
530
531    fn features(&self) -> &Vec<Feature> {
532        // this is not a real implementation, just a placeholder
533        // to satisfy the trait bound
534        static EMPTY: Vec<Feature> = Vec::new();
535        &EMPTY
536    }
537
538    fn llm_records(&self) -> Vec<&LLMRecord> {
539        // this is not a real implementation, just a placeholder
540        // to satisfy the trait bound
541        vec![]
542    }
543}
544
545impl QueueExt for LLMRecord {
546    fn metrics(&self) -> &Vec<Metric> {
547        // this is not a real implementation, just a placeholder
548        // to satisfy the trait bound
549        static EMPTY: Vec<Metric> = Vec::new();
550        &EMPTY
551    }
552
553    fn features(&self) -> &Vec<Feature> {
554        // this is not a real implementation, just a placeholder
555        // to satisfy the trait bound
556        static EMPTY: Vec<Feature> = Vec::new();
557        &EMPTY
558    }
559
560    fn llm_records(&self) -> Vec<&LLMRecord> {
561        vec![self]
562    }
563}
scouter_types/queue/types.rs

scouter_types/queue/
types.rs
