tokenmiser_router/
dsl.rs

1//! Rhai-scripted routing policy.
2//!
3//! A `policy.rhai` file defines `fn route(req)` returning
4//! `#{ provider, model }`. `req` is a `RequestView` exposing `word_count`,
5//! `model`, `tenant`, `has_keyword(s)` and `prompt()`.
6//!
7//! ```rhai
8//! fn route(req) {
9//!     if req.word_count > 500 {
10//!         return #{ provider: "anthropic", model: "claude-opus-4-7" };
11//!     }
12//!     #{ provider: "ollama", model: "ollama:qwen2.5:7b" }
13//! }
14//! ```
15
16use std::path::PathBuf;
17use std::sync::Arc;
18
19use anyhow::{anyhow, Context, Result};
20use parking_lot::RwLock;
21use rhai::{Dynamic, Engine, Map, Scope, AST};
22use serde::{Deserialize, Serialize};
23use tokenmiser_providers::ChatRequest;
24
25use crate::policy::RoutingTarget;
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct RequestView {
29    pub model: String,
30    pub tenant: String,
31    pub word_count: i64,
32    pub prompt: String,
33}
34
35impl RequestView {
36    pub fn from(req: &ChatRequest, tenant: &str) -> Self {
37        let prompt: String = req
38            .messages
39            .iter()
40            .filter(|m| m.role == "user")
41            .filter_map(|m| match &m.content {
42                serde_json::Value::String(s) => Some(s.clone()),
43                _ => None,
44            })
45            .collect::<Vec<_>>()
46            .join("\n");
47        let word_count = prompt.split_whitespace().count() as i64;
48        Self {
49            model: req.model.clone(),
50            tenant: tenant.to_string(),
51            word_count,
52            prompt,
53        }
54    }
55
56    fn into_map(self) -> Map {
57        let mut m = Map::new();
58        m.insert("model".into(), Dynamic::from(self.model));
59        m.insert("tenant".into(), Dynamic::from(self.tenant));
60        m.insert("word_count".into(), Dynamic::from(self.word_count));
61        m.insert("prompt".into(), Dynamic::from(self.prompt));
62        m
63    }
64}
65
66pub struct PolicyEngine {
67    engine: Engine,
68    ast: RwLock<Arc<AST>>,
69    source: RwLock<PathBuf>,
70}
71
72impl PolicyEngine {
73    pub fn load(path: PathBuf) -> Result<Arc<Self>> {
74        let mut engine = Engine::new();
75        engine.set_max_expr_depths(64, 64);
76
77        // Registered as a free function so scripts can call it as
78        // `req.has_keyword("refactor")`.
79        engine.register_fn("has_keyword", |req: Map, kw: &str| -> bool {
80            req.get("prompt")
81                .and_then(|d| d.clone().into_string().ok())
82                .map(|s| s.to_lowercase().contains(&kw.to_lowercase()))
83                .unwrap_or(false)
84        });
85
86        let src = std::fs::read_to_string(&path)
87            .with_context(|| format!("read policy {}", path.display()))?;
88        let ast = engine
89            .compile(&src)
90            .map_err(|e| anyhow!("policy compile: {e}"))?;
91
92        Ok(Arc::new(Self {
93            engine,
94            ast: RwLock::new(Arc::new(ast)),
95            source: RwLock::new(path),
96        }))
97    }
98
99    pub fn reload(&self) -> Result<()> {
100        let path = self.source.read().clone();
101        let src = std::fs::read_to_string(&path)
102            .with_context(|| format!("reload policy {}", path.display()))?;
103        let ast = self
104            .engine
105            .compile(&src)
106            .map_err(|e| anyhow!("policy recompile: {e}"))?;
107        *self.ast.write() = Arc::new(ast);
108        Ok(())
109    }
110
111    pub fn route(&self, req: &ChatRequest, tenant: &str) -> Result<RoutingTarget> {
112        let view = RequestView::from(req, tenant);
113        let mut scope = Scope::new();
114        let ast = self.ast.read().clone();
115        let result: Map = self
116            .engine
117            .call_fn(&mut scope, &ast, "route", (view.into_map(),))
118            .map_err(|e| anyhow!("route() call failed: {e}"))?;
119
120        let provider = result
121            .get("provider")
122            .and_then(|v| v.clone().into_string().ok())
123            .ok_or_else(|| anyhow!("route() result missing `provider`"))?;
124        let model = result
125            .get("model")
126            .and_then(|v| v.clone().into_string().ok())
127            .ok_or_else(|| anyhow!("route() result missing `model`"))?;
128
129        Ok(RoutingTarget { provider, model })
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use std::io::Write;
137    use tokenmiser_providers::ChatMessage;
138
139    fn req(text: &str) -> ChatRequest {
140        ChatRequest {
141            model: "auto".into(),
142            messages: vec![ChatMessage {
143                role: "user".into(),
144                content: serde_json::Value::String(text.into()),
145                extra: Default::default(),
146            }],
147            temperature: None,
148            max_tokens: None,
149            top_p: None,
150            stream: None,
151            extra: Default::default(),
152        }
153    }
154
155    fn write_policy(src: &str) -> PathBuf {
156        let path = std::env::temp_dir().join(format!("tokenmiser-policy-{}.rhai", rand_suffix()));
157        let mut f = std::fs::File::create(&path).unwrap();
158        f.write_all(src.as_bytes()).unwrap();
159        path
160    }
161
162    /// Unique suffix for temp policy files. A timestamp alone collides
163    /// between parallel test threads within one clock tick, letting one test
164    /// truncate another's policy mid-load; the PID separates processes and the
165    /// counter separates threads.
166    fn rand_suffix() -> String {
167        use std::sync::atomic::{AtomicU64, Ordering};
168        static COUNTER: AtomicU64 = AtomicU64::new(0);
169        format!(
170            "{}-{}-{}",
171            std::process::id(),
172            std::time::SystemTime::now()
173                .duration_since(std::time::UNIX_EPOCH)
174                .unwrap()
175                .as_nanos(),
176            COUNTER.fetch_add(1, Ordering::Relaxed),
177        )
178    }
179
180    #[test]
181    fn keyword_rule_routes_to_frontier() {
182        let path = write_policy(
183            r#"
184            fn route(req) {
185                if req.has_keyword("refactor") {
186                    return #{ provider: "anthropic", model: "claude-opus-4-7" };
187                }
188                #{ provider: "ollama", model: "ollama:qwen2.5:7b" }
189            }
190        "#,
191        );
192        let p = PolicyEngine::load(path.clone()).unwrap();
193        let t = p.route(&req("refactor this code"), "t1").unwrap();
194        assert_eq!(t.model, "claude-opus-4-7");
195        let t2 = p.route(&req("what is 2+2"), "t1").unwrap();
196        assert_eq!(t2.model, "ollama:qwen2.5:7b");
197        let _ = std::fs::remove_file(path);
198    }
199
200    #[test]
201    fn word_count_rule_works() {
202        let path = write_policy(
203            r#"
204            fn route(req) {
205                if req.word_count > 50 {
206                    return #{ provider: "anthropic", model: "claude-sonnet-4-6" };
207                }
208                #{ provider: "ollama", model: "ollama:llama2:latest" }
209            }
210        "#,
211        );
212        let p = PolicyEngine::load(path.clone()).unwrap();
213        let long = "word ".repeat(100);
214        assert_eq!(
215            p.route(&req(&long), "t").unwrap().model,
216            "claude-sonnet-4-6"
217        );
218        assert_eq!(
219            p.route(&req("hi"), "t").unwrap().model,
220            "ollama:llama2:latest"
221        );
222        let _ = std::fs::remove_file(path);
223    }
224}
tokenmiser_router/dsl.rs

tokenmiser_router/
dsl.rs
