llmgraph 0.1.1

This library provides a framework for building conversational AI applications with function calling capabilities using a graph-based architecture.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Router agent for conditional routing based on input patterns.

use async_trait::async_trait;
use crate::models::graph::Agent;
use crate::models::tools::ToolRegistryTrait;
use regex::Regex;

/// Configuration for a routing rule
#[derive(Clone)]
pub struct RouteRule {
    /// Pattern to match (regex)
    pub pattern: String,
    /// Node ID to route to if pattern matches
    pub target_node: i32,
    /// Optional description of this route
    pub description: Option<String>,
}

/// Router agent that routes messages based on pattern matching.
/// 
/// # Example
/// ```rust
/// use llmgraph::agents::RouterAgent;
/// use llmgraph::agents::router::RouteRule;
/// 
/// let mut router = RouterAgent::new()
///     .add_route(RouteRule {
///         pattern: r"technical|bug|error".to_string(),
///         target_node: 1,
///         description: Some("Route to technical support".to_string()),
///     })
///     .add_route(RouteRule {
///         pattern: r"sales|pricing|buy".to_string(),
///         target_node: 2,
///         description: Some("Route to sales team".to_string()),
///     })
///     .set_default(3); // Default route
/// ```
pub struct RouterAgent {
    routes: Vec<RouteRule>,
    default_route: Option<i32>,
    case_sensitive: bool,
    name: String,
}

impl RouterAgent {
    /// Create a new router agent
    pub fn new() -> Self {
        Self {
            routes: Vec::new(),
            default_route: None,
            case_sensitive: false,
            name: "Router".to_string(),
        }
    }

    /// Add a routing rule
    pub fn add_route(mut self, rule: RouteRule) -> Self {
        self.routes.push(rule);
        self
    }

    /// Set the default route for unmatched inputs
    pub fn set_default(mut self, node_id: i32) -> Self {
        self.default_route = Some(node_id);
        self
    }

    /// Set whether pattern matching is case-sensitive
    pub fn set_case_sensitive(mut self, sensitive: bool) -> Self {
        self.case_sensitive = sensitive;
        self
    }

    /// Set the agent name
    pub fn with_name(mut self, name: String) -> Self {
        self.name = name;
        self
    }

    /// Find the best matching route for the input
    fn find_route(&self, input: &str) -> Option<i32> {
        let test_input = if self.case_sensitive {
            input.to_string()
        } else {
            input.to_lowercase()
        };

        for rule in &self.routes {
            let pattern = if self.case_sensitive {
                &rule.pattern
            } else {
                &rule.pattern.to_lowercase()
            };

            if let Ok(re) = Regex::new(pattern) {
                if re.is_match(&test_input) {
                    return Some(rule.target_node);
                }
            }
        }

        self.default_route
    }
}

impl Default for RouterAgent {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Agent for RouterAgent {
    async fn run(
        &mut self,
        input: &str,
        _tool_registry: &(dyn ToolRegistryTrait + Send + Sync),
    ) -> (String, Option<i32>) {
        let route = self.find_route(input);
        
        let response = match route {
            Some(node_id) => {
                format!("Routing to node {}: {}", node_id, input)
            }
            None => {
                format!("No matching route found for: {}", input)
            }
        };

        (response, route)
    }

    fn get_name(&self) -> &str {
        &self.name
    }
}