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::ProfileFuncs;
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        ProfileFuncs::__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        ProfileFuncs::__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        ProfileFuncs::__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        ProfileFuncs::__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)| Feature::new(key.extract().unwrap(), value.clone()).unwrap())
239                .collect()
240        } else {
241            Err(TypeError::UnsupportedFeaturesTypeError(
242                features.get_type().name()?.to_string(),
243            ))?
244        };
245        Ok(Features {
246            features,
247            entity_type: EntityType::Feature,
248        })
249    }
250
251    pub fn __str__(&self) -> String {
252        ProfileFuncs::__str__(self)
253    }
254}
255
256impl Features {
257    pub fn iter(&self) -> std::slice::Iter<'_, Feature> {
258        self.features.iter()
259    }
260}
261
262#[pyclass]
263#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
264pub struct FeatureMap {
265    #[pyo3(get)]
266    pub features: HashMap<String, HashMap<String, usize>>,
267}
268
269#[pymethods]
270impl FeatureMap {
271    pub fn __str__(&self) -> String {
272        // serialize the struct to a string
273        ProfileFuncs::__str__(self)
274    }
275}
276
277#[pyclass]
278#[derive(Clone, Serialize, Debug)]
279pub struct Metric {
280    pub name: String,
281    pub value: f64,
282}
283
284#[pymethods]
285impl Metric {
286    #[new]
287    pub fn new(name: String, value: Bound<'_, PyAny>) -> Self {
288        let value = if value.is_instance_of::<PyFloat>() {
289            value.extract::<f64>().unwrap()
290        } else if value.is_instance_of::<PyInt>() {
291            value.extract::<i64>().unwrap() as f64
292        } else {
293            panic!(
294                "Unsupported metric type: {}",
295                value.get_type().name().unwrap()
296            );
297        };
298        let lowercase_name = name.to_lowercase();
299        Metric {
300            name: lowercase_name,
301            value,
302        }
303    }
304
305    pub fn __str__(&self) -> String {
306        ProfileFuncs::__str__(self)
307    }
308}
309
310impl Metric {
311    pub fn new_rs(name: String, value: f64) -> Self {
312        Metric { name, value }
313    }
314}
315
316#[pyclass]
317#[derive(Clone, Serialize, Debug)]
318pub struct Metrics {
319    #[pyo3(get)]
320    pub metrics: Vec<Metric>,
321
322    #[pyo3(get)]
323    pub entity_type: EntityType,
324}
325
326#[pymethods]
327impl Metrics {
328    #[new]
329    pub fn new(metrics: Bound<'_, PyAny>) -> Result<Self, TypeError> {
330        let metrics = if metrics.is_instance_of::<PyList>() {
331            metrics
332                .downcast::<PyList>()
333                .unwrap()
334                .iter()
335                .map(|item| item.extract::<Metric>().unwrap())
336                .collect()
337        } else if metrics.is_instance_of::<PyDict>() {
338            metrics
339                .downcast::<PyDict>()
340                .unwrap()
341                .iter()
342                .map(|(key, value)| Metric::new(key.extract().unwrap(), value))
343                .collect()
344        } else {
345            Err(TypeError::UnsupportedMetricsTypeError(
346                metrics.get_type().name()?.to_string(),
347            ))?
348        };
349        Ok(Metrics {
350            metrics,
351            entity_type: EntityType::Metric,
352        })
353    }
354    pub fn __str__(&self) -> String {
355        ProfileFuncs::__str__(self)
356    }
357}
358
359impl Metrics {
360    pub fn iter(&self) -> std::slice::Iter<'_, Metric> {
361        self.metrics.iter()
362    }
363}
364
365#[pyclass]
366#[derive(Clone, Serialize, Debug)]
367pub struct LLMRecord {
368    pub uid: String,
369
370    pub space: String,
371
372    pub name: String,
373
374    pub version: String,
375
376    pub created_at: DateTime<Utc>,
377
378    pub context: Value,
379
380    pub score: Value,
381
382    pub prompt: Option<Value>,
383
384    #[pyo3(get)]
385    pub entity_type: EntityType,
386}
387
388#[pymethods]
389impl LLMRecord {
390    #[new]
391    #[pyo3(signature = (
392        context,
393        prompt=None,
394    ))]
395
396    /// Creates a new LLMRecord instance.
397    /// The context is either a python dictionary or a pydantic basemodel.
398    pub fn new(
399        py: Python<'_>,
400        context: Bound<'_, PyAny>,
401        prompt: Option<Bound<'_, PyAny>>,
402    ) -> Result<Self, TypeError> {
403        // check if context is a PyDict or PyObject(Pydantic model)
404        let context_val = if context.is_instance_of::<PyDict>() {
405            pyobject_to_json(&context)?
406        } else if is_pydantic_model(py, &context)? {
407            // Dump pydantic model to dictionary
408            let model = context.call_method0("model_dump")?;
409
410            // Serialize the dictionary to JSON
411            pyobject_to_json(&model)?
412        } else {
413            Err(TypeError::MustBeDictOrBaseModel)?
414        };
415
416        let prompt: Option<Value> = match prompt {
417            Some(p) => {
418                if p.is_instance_of::<Prompt>() {
419                    let prompt = p.extract::<Prompt>()?;
420                    Some(serde_json::to_value(prompt)?)
421                } else {
422                    Some(pyobject_to_json(&p)?)
423                }
424            }
425            None => None,
426        };
427
428        Ok(LLMRecord {
429            uid: create_uuid7(),
430            created_at: Utc::now(),
431            space: String::new(),
432            name: String::new(),
433            version: String::new(),
434            context: context_val,
435            score: Value::Null,
436            prompt,
437            entity_type: EntityType::LLM,
438        })
439    }
440
441    #[getter]
442    pub fn context<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
443        Ok(json_to_pyobject_value(py, &self.context)?
444            .into_bound_py_any(py)?
445            .clone())
446    }
447}
448
449impl LLMRecord {
450    pub fn new_rs(context: Option<Value>, prompt: Option<Value>) -> Self {
451        LLMRecord {
452            context: context.unwrap_or(Value::Object(serde_json::Map::new())),
453            prompt,
454            entity_type: EntityType::LLM,
455            uid: create_uuid7(),
456            created_at: Utc::now(),
457            space: String::new(),
458            name: String::new(),
459            version: String::new(),
460            score: Value::Null,
461        }
462    }
463
464    pub fn __str__(&self) -> String {
465        ProfileFuncs::__str__(self)
466    }
467}
468
469#[derive(Debug)]
470pub enum QueueItem {
471    Features(Features),
472    Metrics(Metrics),
473    LLM(Box<LLMRecord>),
474}
475
476impl QueueItem {
477    /// Helper for extracting an Entity from a Python object
478    pub fn from_py_entity(entity: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
479        let entity_type = entity.getattr("entity_type")?.extract::<EntityType>()?;
480
481        match entity_type {
482            EntityType::Feature => {
483                let features = entity.extract::<Features>()?;
484                Ok(QueueItem::Features(features))
485            }
486            EntityType::Metric => {
487                let metrics = entity.extract::<Metrics>()?;
488                Ok(QueueItem::Metrics(metrics))
489            }
490            EntityType::LLM => {
491                // LLM is not supported in this context
492                let llm = entity.extract::<LLMRecord>()?;
493                Ok(QueueItem::LLM(Box::new(llm)))
494            }
495        }
496    }
497}
498
499pub trait QueueExt: Send + Sync {
500    fn metrics(&self) -> &Vec<Metric>;
501    fn features(&self) -> &Vec<Feature>;
502    fn llm_records(&self) -> Vec<&LLMRecord>;
503}
504
505impl QueueExt for Features {
506    fn metrics(&self) -> &Vec<Metric> {
507        // this is not a real implementation, just a placeholder
508        // to satisfy the trait bound
509        static EMPTY: Vec<Metric> = Vec::new();
510        &EMPTY
511    }
512
513    fn features(&self) -> &Vec<Feature> {
514        &self.features
515    }
516
517    fn llm_records(&self) -> Vec<&LLMRecord> {
518        // this is not a real implementation, just a placeholder
519        // to satisfy the trait bound
520        vec![]
521    }
522}
523
524impl QueueExt for Metrics {
525    fn metrics(&self) -> &Vec<Metric> {
526        &self.metrics
527    }
528
529    fn features(&self) -> &Vec<Feature> {
530        // this is not a real implementation, just a placeholder
531        // to satisfy the trait bound
532        static EMPTY: Vec<Feature> = Vec::new();
533        &EMPTY
534    }
535
536    fn llm_records(&self) -> Vec<&LLMRecord> {
537        // this is not a real implementation, just a placeholder
538        // to satisfy the trait bound
539        vec![]
540    }
541}
542
543impl QueueExt for LLMRecord {
544    fn metrics(&self) -> &Vec<Metric> {
545        // this is not a real implementation, just a placeholder
546        // to satisfy the trait bound
547        static EMPTY: Vec<Metric> = Vec::new();
548        &EMPTY
549    }
550
551    fn features(&self) -> &Vec<Feature> {
552        // this is not a real implementation, just a placeholder
553        // to satisfy the trait bound
554        static EMPTY: Vec<Feature> = Vec::new();
555        &EMPTY
556    }
557
558    fn llm_records(&self) -> Vec<&LLMRecord> {
559        vec![self]
560    }
561}
scouter_types/queue/types.rs

scouter_types/queue/
types.rs
