scouter_types/queue/
types.rs

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

scouter_types/queue/
types.rs
