1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use chrono::{Datelike, Local, Timelike, Utc};
6use serde::Deserialize;
7use serde_json::Value;
8
9use ai_agents_core::{ChatMessage, LLMProvider, Result};
10use ai_agents_state::{CompareOp, ContextMatcher, StateMatcher, TimeMatcher, ToolCondition};
11
12#[derive(Debug, Clone)]
13pub struct ToolCallRecord {
14 pub tool_id: String,
15 pub result: Value,
16 pub timestamp: chrono::DateTime<chrono::Utc>,
17}
18
19#[derive(Debug, Clone, Default)]
20pub struct EvaluationContext {
21 pub context: HashMap<String, Value>,
22 pub state_name: Option<String>,
23 pub state_turn_count: u32,
24 pub previous_state: Option<String>,
25 pub called_tools: Vec<ToolCallRecord>,
26 pub recent_messages: Vec<ChatMessage>,
27}
28
29impl EvaluationContext {
30 pub fn new() -> Self {
31 Self::default()
32 }
33
34 pub fn with_context(mut self, context: HashMap<String, Value>) -> Self {
35 self.context = context;
36 self
37 }
38
39 pub fn with_state(
40 mut self,
41 name: Option<String>,
42 turn_count: u32,
43 previous: Option<String>,
44 ) -> Self {
45 self.state_name = name;
46 self.state_turn_count = turn_count;
47 self.previous_state = previous;
48 self
49 }
50
51 pub fn with_called_tools(mut self, tools: Vec<ToolCallRecord>) -> Self {
52 self.called_tools = tools;
53 self
54 }
55
56 pub fn with_messages(mut self, messages: Vec<ChatMessage>) -> Self {
57 self.recent_messages = messages;
58 self
59 }
60}
61
62#[async_trait]
63pub trait LLMGetter: Send + Sync {
64 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>>;
65}
66
67pub struct ConditionEvaluator<G: LLMGetter> {
68 llm_getter: G,
69}
70
71impl<G: LLMGetter> ConditionEvaluator<G> {
72 pub fn new(llm_getter: G) -> Self {
73 Self { llm_getter }
74 }
75
76 pub async fn evaluate(
77 &self,
78 condition: &ToolCondition,
79 ctx: &EvaluationContext,
80 ) -> Result<bool> {
81 match condition {
82 ToolCondition::Context(matchers) => Ok(self.evaluate_context(matchers, &ctx.context)),
83 ToolCondition::State(matcher) => Ok(self.evaluate_state(matcher, ctx)),
84 ToolCondition::AfterTool(tool_id) => {
85 Ok(ctx.called_tools.iter().any(|t| &t.tool_id == tool_id))
86 }
87 ToolCondition::ToolResult { tool, result } => {
88 Ok(self.evaluate_tool_result(tool, result, &ctx.called_tools))
89 }
90 ToolCondition::Semantic {
91 when,
92 llm,
93 threshold,
94 } => self.evaluate_semantic(when, llm, *threshold, ctx).await,
95 ToolCondition::Time(matcher) => Ok(self.evaluate_time(matcher)),
96 ToolCondition::All(conditions) => {
97 for cond in conditions {
98 if !Box::pin(self.evaluate(cond, ctx)).await? {
99 return Ok(false);
100 }
101 }
102 Ok(true)
103 }
104 ToolCondition::Any(conditions) => {
105 for cond in conditions {
106 if Box::pin(self.evaluate(cond, ctx)).await? {
107 return Ok(true);
108 }
109 }
110 Ok(false)
111 }
112 ToolCondition::Not(inner) => Ok(!Box::pin(self.evaluate(inner, ctx)).await?),
113 }
114 }
115
116 fn evaluate_context(
117 &self,
118 matchers: &HashMap<String, ContextMatcher>,
119 context: &HashMap<String, Value>,
120 ) -> bool {
121 for (path, matcher) in matchers {
122 let value = self.get_context_value(path, context);
123 if !self.match_value(value.as_ref(), matcher) {
124 return false;
125 }
126 }
127 true
128 }
129
130 fn get_context_value(&self, path: &str, context: &HashMap<String, Value>) -> Option<Value> {
131 ai_agents_core::get_dot_path_from_map(context, path)
132 }
133
134 fn match_value(&self, value: Option<&Value>, matcher: &ContextMatcher) -> bool {
135 match matcher {
136 ContextMatcher::Exact(expected) => value.map(|v| v == expected).unwrap_or(false),
137 ContextMatcher::Exists { exists } => {
138 let has_value = value.is_some() && value != Some(&Value::Null);
139 *exists == has_value
140 }
141 ContextMatcher::Compare(op) => {
142 let Some(val) = value else {
143 return false;
144 };
145 self.compare_value(val, op)
146 }
147 }
148 }
149
150 fn compare_value(&self, value: &Value, op: &CompareOp) -> bool {
151 match op {
152 CompareOp::Eq(expected) => value == expected,
153 CompareOp::Neq(expected) => value != expected,
154 CompareOp::Gt(n) => value.as_f64().map(|v| v > *n).unwrap_or(false),
155 CompareOp::Gte(n) => value.as_f64().map(|v| v >= *n).unwrap_or(false),
156 CompareOp::Lt(n) => value.as_f64().map(|v| v < *n).unwrap_or(false),
157 CompareOp::Lte(n) => value.as_f64().map(|v| v <= *n).unwrap_or(false),
158 CompareOp::In(values) => values.contains(value),
159 CompareOp::Contains(s) => value
160 .as_str()
161 .map(|v| v.contains(s))
162 .or_else(|| {
163 value
164 .as_array()
165 .map(|arr| arr.iter().any(|v| v.as_str() == Some(s)))
166 })
167 .unwrap_or(false),
168 }
169 }
170
171 fn evaluate_state(&self, matcher: &StateMatcher, ctx: &EvaluationContext) -> bool {
172 if let Some(ref expected_name) = matcher.name
173 && ctx.state_name.as_ref() != Some(expected_name)
174 {
175 return false;
176 }
177
178 if let Some(ref turn_op) = matcher.turn_count {
179 let turn_count = ctx.state_turn_count as f64;
180 if !self.compare_value(
181 &Value::Number(serde_json::Number::from_f64(turn_count).unwrap_or(0.into())),
182 turn_op,
183 ) {
184 return false;
185 }
186 }
187
188 if let Some(ref expected_prev) = matcher.previous
189 && ctx.previous_state.as_ref() != Some(expected_prev)
190 {
191 return false;
192 }
193
194 true
195 }
196
197 fn evaluate_tool_result(
198 &self,
199 tool: &str,
200 expected: &HashMap<String, Value>,
201 called_tools: &[ToolCallRecord],
202 ) -> bool {
203 let tool_record = called_tools.iter().rev().find(|t| t.tool_id == tool);
204
205 let Some(record) = tool_record else {
206 return false;
207 };
208
209 let result_obj = match &record.result {
210 Value::Object(obj) => obj,
211 _ => return false,
212 };
213
214 for (key, expected_value) in expected {
215 match result_obj.get(key) {
216 Some(actual) if actual == expected_value => continue,
217 _ => return false,
218 }
219 }
220
221 true
222 }
223
224 fn evaluate_time(&self, matcher: &TimeMatcher) -> bool {
225 let now = if let Some(ref tz) = matcher.timezone {
226 if tz == "utc" || tz == "UTC" {
227 Utc::now().with_timezone(&Utc).naive_local()
228 } else {
229 Local::now().naive_local()
230 }
231 } else {
232 Local::now().naive_local()
233 };
234
235 if let Some(ref hours_op) = matcher.hours {
236 let hour = now.hour() as f64;
237 if !self.compare_value(&serde_json::json!(hour), hours_op) {
238 return false;
239 }
240 }
241
242 if let Some(ref days) = matcher.day_of_week {
243 let day_name = match now.weekday() {
244 chrono::Weekday::Mon => "monday",
245 chrono::Weekday::Tue => "tuesday",
246 chrono::Weekday::Wed => "wednesday",
247 chrono::Weekday::Thu => "thursday",
248 chrono::Weekday::Fri => "friday",
249 chrono::Weekday::Sat => "saturday",
250 chrono::Weekday::Sun => "sunday",
251 };
252
253 if !days.iter().any(|d| d.to_lowercase() == day_name) {
254 return false;
255 }
256 }
257
258 true
259 }
260
261 async fn evaluate_semantic(
262 &self,
263 condition: &str,
264 llm_alias: &str,
265 threshold: f32,
266 ctx: &EvaluationContext,
267 ) -> Result<bool> {
268 let llm = match self.llm_getter.get_llm(llm_alias) {
269 Some(l) => l,
270 None => {
271 tracing::warn!(llm = llm_alias, "LLM not found for semantic evaluation");
272 return Ok(false);
273 }
274 };
275
276 let conversation_summary = ctx
277 .recent_messages
278 .iter()
279 .take(10)
280 .map(|m| format!("{:?}: {}", m.role, m.content))
281 .collect::<Vec<_>>()
282 .join("\n");
283
284 let prompt = format!(
285 r#"Based on the conversation below, evaluate if this condition is TRUE or FALSE.
286
287Condition to evaluate: "{}"
288
289Recent conversation:
290{}
291
292Respond with ONLY a JSON object:
293{{"result": true, "confidence": 0.9, "reason": "brief explanation"}}
294or
295{{"result": false, "confidence": 0.9, "reason": "brief explanation"}}"#,
296 condition, conversation_summary
297 );
298
299 let messages = vec![ChatMessage::user(&prompt)];
300 let response = llm.complete(&messages, None).await?;
301
302 let parsed: SemanticEvalResult =
303 serde_json::from_str(&response.content).unwrap_or(SemanticEvalResult {
304 result: false,
305 confidence: 0.0,
306 reason: "Failed to parse response".to_string(),
307 });
308
309 tracing::debug!(
310 condition = condition,
311 result = parsed.result,
312 confidence = parsed.confidence,
313 threshold = threshold,
314 reason = %parsed.reason,
315 "Semantic evaluation"
316 );
317
318 Ok(parsed.result && parsed.confidence >= threshold)
319 }
320}
321
322#[derive(Debug, Deserialize)]
323struct SemanticEvalResult {
324 result: bool,
325 confidence: f32,
326 reason: String,
327}
328
329pub struct SimpleLLMGetter {
330 llms: HashMap<String, Arc<dyn LLMProvider>>,
331}
332
333impl SimpleLLMGetter {
334 pub fn new() -> Self {
335 Self {
336 llms: HashMap::new(),
337 }
338 }
339
340 pub fn with_llm(mut self, alias: &str, llm: Arc<dyn LLMProvider>) -> Self {
341 self.llms.insert(alias.to_string(), llm);
342 self
343 }
344}
345
346impl Default for SimpleLLMGetter {
347 fn default() -> Self {
348 Self::new()
349 }
350}
351
352impl LLMGetter for SimpleLLMGetter {
353 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
354 self.llms.get(alias).cloned()
355 }
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361
362 struct NoOpLLMGetter;
363
364 impl LLMGetter for NoOpLLMGetter {
365 fn get_llm(&self, _alias: &str) -> Option<Arc<dyn LLMProvider>> {
366 None
367 }
368 }
369
370 #[tokio::test]
371 async fn test_context_condition_exact() {
372 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
373
374 let mut context = HashMap::new();
375 context.insert(
376 "user".to_string(),
377 serde_json::json!({
378 "verified": true,
379 "tier": "premium"
380 }),
381 );
382
383 let ctx = EvaluationContext::new().with_context(context);
384
385 let mut matchers = HashMap::new();
386 matchers.insert(
387 "user.verified".to_string(),
388 ContextMatcher::Exact(Value::Bool(true)),
389 );
390
391 let condition = ToolCondition::Context(matchers);
392 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
393 }
394
395 #[tokio::test]
396 async fn test_context_condition_exists() {
397 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
398
399 let mut context = HashMap::new();
400 context.insert("name".to_string(), Value::String("Alice".into()));
401
402 let ctx = EvaluationContext::new().with_context(context);
403
404 let mut matchers = HashMap::new();
405 matchers.insert("name".to_string(), ContextMatcher::Exists { exists: true });
406 matchers.insert(
407 "email".to_string(),
408 ContextMatcher::Exists { exists: false },
409 );
410
411 let condition = ToolCondition::Context(matchers);
412 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
413 }
414
415 #[tokio::test]
416 async fn test_context_condition_compare() {
417 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
418
419 let mut context = HashMap::new();
420 context.insert("balance".to_string(), serde_json::json!(150.0));
421
422 let ctx = EvaluationContext::new().with_context(context);
423
424 let mut matchers = HashMap::new();
425 matchers.insert(
426 "balance".to_string(),
427 ContextMatcher::Compare(CompareOp::Gte(100.0)),
428 );
429
430 let condition = ToolCondition::Context(matchers);
431 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
432 }
433
434 #[tokio::test]
435 async fn test_state_condition() {
436 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
437
438 let ctx = EvaluationContext::new().with_state(
439 Some("checkout".to_string()),
440 5,
441 Some("browsing".to_string()),
442 );
443
444 let condition = ToolCondition::State(StateMatcher {
445 name: Some("checkout".to_string()),
446 turn_count: Some(CompareOp::Gte(3.0)),
447 previous: Some("browsing".to_string()),
448 });
449
450 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
451 }
452
453 #[tokio::test]
454 async fn test_after_tool_condition() {
455 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
456
457 let ctx = EvaluationContext::new().with_called_tools(vec![ToolCallRecord {
458 tool_id: "search".to_string(),
459 result: serde_json::json!({"found": true}),
460 timestamp: Utc::now(),
461 }]);
462
463 let condition = ToolCondition::AfterTool("search".to_string());
464 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
465
466 let condition2 = ToolCondition::AfterTool("calculate".to_string());
467 assert!(!evaluator.evaluate(&condition2, &ctx).await.unwrap());
468 }
469
470 #[tokio::test]
471 async fn test_tool_result_condition() {
472 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
473
474 let ctx = EvaluationContext::new().with_called_tools(vec![ToolCallRecord {
475 tool_id: "verify_purchase".to_string(),
476 result: serde_json::json!({
477 "valid": true,
478 "refundable": true
479 }),
480 timestamp: Utc::now(),
481 }]);
482
483 let mut expected = HashMap::new();
484 expected.insert("valid".to_string(), Value::Bool(true));
485 expected.insert("refundable".to_string(), Value::Bool(true));
486
487 let condition = ToolCondition::ToolResult {
488 tool: "verify_purchase".to_string(),
489 result: expected,
490 };
491
492 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
493 }
494
495 #[tokio::test]
496 async fn test_all_condition() {
497 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
498
499 let mut context = HashMap::new();
500 context.insert("verified".to_string(), Value::Bool(true));
501 context.insert("balance".to_string(), serde_json::json!(100.0));
502
503 let ctx = EvaluationContext::new().with_context(context);
504
505 let mut m1 = HashMap::new();
506 m1.insert(
507 "verified".to_string(),
508 ContextMatcher::Exact(Value::Bool(true)),
509 );
510
511 let mut m2 = HashMap::new();
512 m2.insert(
513 "balance".to_string(),
514 ContextMatcher::Compare(CompareOp::Gte(50.0)),
515 );
516
517 let condition =
518 ToolCondition::All(vec![ToolCondition::Context(m1), ToolCondition::Context(m2)]);
519
520 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
521 }
522
523 #[tokio::test]
524 async fn test_any_condition() {
525 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
526
527 let mut context = HashMap::new();
528 context.insert("tier".to_string(), Value::String("basic".into()));
529
530 let ctx = EvaluationContext::new().with_context(context);
531
532 let mut m1 = HashMap::new();
533 m1.insert(
534 "tier".to_string(),
535 ContextMatcher::Exact(Value::String("premium".into())),
536 );
537
538 let mut m2 = HashMap::new();
539 m2.insert(
540 "tier".to_string(),
541 ContextMatcher::Exact(Value::String("basic".into())),
542 );
543
544 let condition =
545 ToolCondition::Any(vec![ToolCondition::Context(m1), ToolCondition::Context(m2)]);
546
547 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
548 }
549
550 #[tokio::test]
551 async fn test_not_condition() {
552 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
553
554 let mut context = HashMap::new();
555 context.insert("blocked".to_string(), Value::Bool(false));
556
557 let ctx = EvaluationContext::new().with_context(context);
558
559 let mut matchers = HashMap::new();
560 matchers.insert(
561 "blocked".to_string(),
562 ContextMatcher::Exact(Value::Bool(true)),
563 );
564
565 let condition = ToolCondition::Not(Box::new(ToolCondition::Context(matchers)));
566 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
567 }
568
569 #[tokio::test]
570 async fn test_time_condition_day_of_week() {
571 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
572 let ctx = EvaluationContext::new();
573
574 let all_days = vec![
575 "monday".to_string(),
576 "tuesday".to_string(),
577 "wednesday".to_string(),
578 "thursday".to_string(),
579 "friday".to_string(),
580 "saturday".to_string(),
581 "sunday".to_string(),
582 ];
583
584 let condition = ToolCondition::Time(TimeMatcher {
585 hours: None,
586 day_of_week: Some(all_days),
587 timezone: None,
588 });
589
590 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
591 }
592
593 #[tokio::test]
594 async fn test_compare_in() {
595 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
596
597 let mut context = HashMap::new();
598 context.insert("status".to_string(), Value::String("active".into()));
599
600 let ctx = EvaluationContext::new().with_context(context);
601
602 let mut matchers = HashMap::new();
603 matchers.insert(
604 "status".to_string(),
605 ContextMatcher::Compare(CompareOp::In(vec![
606 Value::String("active".into()),
607 Value::String("pending".into()),
608 ])),
609 );
610
611 let condition = ToolCondition::Context(matchers);
612 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
613 }
614
615 #[tokio::test]
616 async fn test_compare_contains_string() {
617 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
618
619 let mut context = HashMap::new();
620 context.insert(
621 "email".to_string(),
622 Value::String("user@example.com".into()),
623 );
624
625 let ctx = EvaluationContext::new().with_context(context);
626
627 let mut matchers = HashMap::new();
628 matchers.insert(
629 "email".to_string(),
630 ContextMatcher::Compare(CompareOp::Contains("@example.com".into())),
631 );
632
633 let condition = ToolCondition::Context(matchers);
634 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
635 }
636
637 #[tokio::test]
638 async fn test_nested_context_path() {
639 let evaluator = ConditionEvaluator::new(NoOpLLMGetter);
640
641 let mut context = HashMap::new();
642 context.insert(
643 "user".to_string(),
644 serde_json::json!({
645 "profile": {
646 "settings": {
647 "notifications": true
648 }
649 }
650 }),
651 );
652
653 let ctx = EvaluationContext::new().with_context(context);
654
655 let mut matchers = HashMap::new();
656 matchers.insert(
657 "user.profile.settings.notifications".to_string(),
658 ContextMatcher::Exact(Value::Bool(true)),
659 );
660
661 let condition = ToolCondition::Context(matchers);
662 assert!(evaluator.evaluate(&condition, &ctx).await.unwrap());
663 }
664
665 #[test]
666 fn test_simple_llm_getter() {
667 let getter = SimpleLLMGetter::new();
668 assert!(getter.get_llm("nonexistent").is_none());
669 }
670}