Skip to main content

scouter_evaluate/tasks/
trace.rs

1use crate::evaluate::trace::TraceContextBuilder;
2use crate::tasks::evaluator::AssertionEvaluator;
3use crate::{error::EvaluationError, tasks::traits::EvaluationTask};
4use pyo3::pyfunction;
5use scouter_types::genai::{AssertionResult, AssertionResults, TraceAssertionTask};
6use scouter_types::sql::TraceSpan;
7use serde_json::Value;
8use std::collections::HashMap;
9use std::sync::Arc;
10use tracing::error;
11
12impl EvaluationTask for TraceAssertionTask {
13    fn execute(&self, context: &Value) -> Result<AssertionResult, EvaluationError> {
14        AssertionEvaluator::evaluate_assertion(context, self)
15    }
16}
17
18#[pyfunction]
19#[pyo3(signature = (tasks, spans))]
20pub fn execute_trace_assertion_tasks(
21    tasks: Vec<TraceAssertionTask>,
22    spans: Vec<TraceSpan>,
23) -> Result<AssertionResults, EvaluationError> {
24    let context_builder = TraceContextBuilder::new(Arc::new(spans));
25
26    let results: HashMap<String, AssertionResult> = tasks
27        .iter()
28        .map(|task| {
29            let context = context_builder.build_context(&task.assertion)?;
30            task.execute(&context)
31                .map(|result| (task.id.clone(), result))
32        })
33        .collect::<Result<HashMap<String, AssertionResult>, EvaluationError>>()?;
34
35    Ok(AssertionResults { results })
36}
37
38pub(crate) fn execute_trace_assertions(
39    context: &TraceContextBuilder,
40    tasks: &[TraceAssertionTask],
41) -> Result<AssertionResults, EvaluationError> {
42    if tasks.is_empty() {
43        return Ok(AssertionResults {
44            results: HashMap::new(),
45        });
46    }
47
48    // If there are tasks but no spans, fail all assertions
49    if context.spans.is_empty() {
50        let results: HashMap<String, AssertionResult> = tasks
51            .iter()
52            .map(|task| {
53                (
54                    task.id.clone(),
55                    AssertionResult {
56                        passed: false,
57                        actual: serde_json::Value::Null,
58                        expected: serde_json::Value::Null,
59                        message: "No spans available for evaluation".to_string(),
60                    },
61                )
62            })
63            .collect();
64
65        return Ok(AssertionResults { results });
66    }
67
68    let results: HashMap<String, AssertionResult> = tasks
69        .iter()
70        .map(|task| {
71            let context = context.build_context(&task.assertion)?;
72            task.execute(&context)
73                .map(|result| (task.id.clone(), result))
74        })
75        .collect::<Result<HashMap<String, AssertionResult>, EvaluationError>>()
76        .inspect_err(|e| {
77            error!("Error executing trace assertions: {:?}", e);
78        })?;
79
80    Ok(AssertionResults { results })
81}