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

scouter_types/queue/
types.rs
