scouter_evaluate/tasks/
trace.rs1use 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 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}