laya/
router.rs

1//! Route a request to the checkpoint best suited to it.
2//!
3//! The English checkpoint does not gently degrade off English: on 20-option
4//! MASSIVE intent it scores 0.100 on Hindi and 0.103 on Korean against 0.050
5//! for random guessing, and reports high confidence while doing so. Script
6//! detection is therefore the primary routing signal, and the Latin-script
7//! language guess is secondary.
8//!
9//! `typed-decisions` is never selected automatically unless you opt in with
10//! [`Router::auto_task_detection`] or ask for it by name: it is fine-tuned on
11//! four specific synthetic workflows and should not be a silent default.
12
13use laya_core::lang;
14use serde::Serialize;
15use serde_json::{Map, Value};
16use std::collections::HashMap;
17use std::sync::Arc;
18
19use crate::agent::Agent;
20use crate::{Error, Result};
21
22/// The bundle repository carrying all three checkpoints.
23pub const BUNDLE_REPO: &str = "convaiinnovations/laya";
24
25/// Which checkpoint to use.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
27pub enum ModelName {
28    English,
29    Multilingual,
30    TypedDecisions,
31}
32
33impl ModelName {
34    pub fn as_str(self) -> &'static str {
35        match self {
36            Self::English => "english",
37            Self::Multilingual => "multilingual",
38            Self::TypedDecisions => "typed-decisions",
39        }
40    }
41
42    /// Repository and optional subfolder inside the bundle repository.
43    pub fn bundle_location(self) -> (&'static str, Option<&'static str>) {
44        match self {
45            Self::English => (BUNDLE_REPO, None),
46            Self::Multilingual => (BUNDLE_REPO, Some("multilingual")),
47            Self::TypedDecisions => (BUNDLE_REPO, Some("typed-decisions")),
48        }
49    }
50
51    /// The standalone repository for this checkpoint.
52    pub fn standalone_repo(self) -> &'static str {
53        match self {
54            Self::English => "convaiinnovations/laya",
55            Self::Multilingual => "convaiinnovations/laya-multilingual",
56            Self::TypedDecisions => "convaiinnovations/laya-typed-decisions",
57        }
58    }
59
60    /// Resolve a name or alias, accepting everything upstream accepts.
61    pub fn parse(name: &str) -> Result<Self> {
62        let key = name.trim().to_lowercase();
63        match key.as_str() {
64            "english" | "en" | "eng" | "laya" | "default" => Ok(Self::English),
65            "multilingual" | "multi" | "ml" | "laya-multilingual" => Ok(Self::Multilingual),
66            "typed-decisions" | "typed" | "typed_decisions" | "laya-typed-decisions"
67            | "decisions" => Ok(Self::TypedDecisions),
68            _ => Err(Error::Config(format!(
69                "unknown model {name:?}; choose one of [\"english\", \"multilingual\", \"typed-decisions\"]"
70            ))),
71        }
72    }
73}
74
75/// Question-id signatures of the four typed-decisions workflows.
76const TYPED_DECISION_WORKFLOWS: [(&str, &[&str]); 4] = [
77    (
78        "agent_trace_observability",
79        &["action", "needs_review", "outcome", "risk", "urgency"],
80    ),
81    (
82        "customer_service",
83        &["action", "category", "churn_risk", "needs_human", "urgency"],
84    ),
85    (
86        "invoice_processing",
87        &[
88            "discrepancy_severity",
89            "disposition",
90            "duplicate",
91            "matches_order",
92            "urgency",
93        ],
94    ),
95    (
96        "security_incidents",
97        &[
98            "credential_compromise",
99            "disposition",
100            "severity",
101            "true_positive",
102            "urgency",
103        ],
104    ),
105];
106
107/// Name of the typed-decisions workflow whose question ids these are.
108///
109/// Requires an exact id-set match, so an unrelated schema that happens to
110/// contain `urgency` is never captured.
111pub fn match_typed_decisions_workflow(questions: &Map<String, Value>) -> Option<&'static str> {
112    let ids: std::collections::BTreeSet<&str> = questions.keys().map(String::as_str).collect();
113    TYPED_DECISION_WORKFLOWS
114        .iter()
115        .find(|(_, signature)| {
116            signature.len() == ids.len() && signature.iter().all(|id| ids.contains(id))
117        })
118        .map(|(name, _)| *name)
119}
120
121/// The routing outcome: which model, why, and what was detected.
122#[derive(Clone, Debug, Serialize)]
123pub struct RouteDecision {
124    pub model: String,
125    pub repo: String,
126    pub reason: String,
127    pub detection: Option<lang::Detection>,
128    pub workflow: Option<String>,
129}
130
131impl RouteDecision {
132    pub fn to_value(&self) -> Value {
133        serde_json::to_value(self).unwrap_or(Value::Null)
134    }
135}
136
137/// Lazily loads checkpoints and sends each request to the right one.
138pub struct Router {
139    standalone_repos: bool,
140    default: ModelName,
141    auto_task_detection: bool,
142    max_loaded: usize,
143    agents: HashMap<ModelName, Arc<Agent>>,
144    order: Vec<ModelName>,
145}
146
147impl Default for Router {
148    fn default() -> Self {
149        Self::new()
150    }
151}
152
153impl Router {
154    pub fn new() -> Self {
155        Self {
156            standalone_repos: false,
157            default: ModelName::English,
158            auto_task_detection: false,
159            max_loaded: 1,
160            agents: HashMap::new(),
161            order: Vec::new(),
162        }
163    }
164
165    /// Use each checkpoint's own repository rather than the bundle.
166    pub fn standalone_repos(mut self, standalone: bool) -> Self {
167        self.standalone_repos = standalone;
168        self
169    }
170
171    /// Checkpoint used when detection is inconclusive.
172    pub fn default_model(mut self, model: ModelName) -> Self {
173        self.default = model;
174        self
175    }
176
177    /// Allow the typed-decisions checkpoint to be selected by workflow match.
178    pub fn auto_task_detection(mut self, enabled: bool) -> Self {
179        self.auto_task_detection = enabled;
180        self
181    }
182
183    /// How many checkpoints stay resident; least-recently-used is evicted.
184    pub fn max_loaded(mut self, max_loaded: usize) -> Self {
185        self.max_loaded = max_loaded.max(1);
186        self
187    }
188
189    fn repo_string(&self, model: ModelName) -> String {
190        if self.standalone_repos {
191            model.standalone_repo().to_string()
192        } else {
193            match model.bundle_location() {
194                (repo, Some(sub)) => format!("{repo}/{sub}"),
195                (repo, None) => repo.to_string(),
196            }
197        }
198    }
199
200    /// Register an already-built agent rather than loading a second copy.
201    pub fn attach(&mut self, model: ModelName, agent: Arc<Agent>) {
202        self.agents.insert(model, agent);
203        self.touch(model);
204        self.max_loaded = self.max_loaded.max(self.agents.len());
205    }
206
207    fn touch(&mut self, model: ModelName) {
208        self.order.retain(|entry| *entry != model);
209        self.order.push(model);
210    }
211
212    /// Checkpoints currently resident, least-recently-used first.
213    pub fn loaded(&self) -> &[ModelName] {
214        &self.order
215    }
216
217    /// Raise the residency cap so a preload is not immediately evicted.
218    pub(crate) fn set_max_loaded_at_least(&mut self, count: usize) {
219        self.max_loaded = self.max_loaded.max(count).max(1);
220    }
221
222    /// Free one checkpoint, or all of them.
223    pub fn unload(&mut self, model: Option<ModelName>) {
224        match model {
225            Some(model) => {
226                self.agents.remove(&model);
227                self.order.retain(|entry| *entry != model);
228            }
229            None => {
230                self.agents.clear();
231                self.order.clear();
232            }
233        }
234    }
235
236    /// Fetch an attached agent, if present.
237    pub fn get(&mut self, model: ModelName) -> Option<Arc<Agent>> {
238        let agent = self.agents.get(&model).cloned();
239        if agent.is_some() {
240            self.touch(model);
241        }
242        agent
243    }
244
245    /// Decide which checkpoint to use, without loading or running anything.
246    ///
247    /// Precedence: explicit `model` > explicit `task` > detected workflow
248    /// (opt-in) > explicit `lang` > detected script/language > default.
249    pub fn route(
250        &self,
251        state: &Value,
252        questions: Option<&Map<String, Value>>,
253        model: Option<&str>,
254        task: Option<&str>,
255        language: Option<&str>,
256    ) -> Result<RouteDecision> {
257        let empty = Map::new();
258        let questions = questions.unwrap_or(&empty);
259
260        if let Some(model) = model {
261            let chosen = ModelName::parse(model)?;
262            return Ok(RouteDecision {
263                model: chosen.as_str().into(),
264                repo: self.repo_string(chosen),
265                reason: format!("explicit model={model:?}"),
266                detection: None,
267                workflow: None,
268            });
269        }
270
271        if let Some(task) = task {
272            let normalised = task.to_lowercase().replace('-', "_");
273            let chosen = if normalised == "typed_decisions" {
274                ModelName::TypedDecisions
275            } else {
276                ModelName::parse(task)?
277            };
278            return Ok(RouteDecision {
279                model: chosen.as_str().into(),
280                repo: self.repo_string(chosen),
281                reason: format!("explicit task={task:?}"),
282                detection: None,
283                workflow: None,
284            });
285        }
286
287        let workflow = match_typed_decisions_workflow(questions);
288        if let (Some(workflow), true) = (workflow, self.auto_task_detection) {
289            return Ok(RouteDecision {
290                model: ModelName::TypedDecisions.as_str().into(),
291                repo: self.repo_string(ModelName::TypedDecisions),
292                reason: format!("question ids match the {workflow:?} typed-decisions workflow"),
293                detection: None,
294                workflow: Some(workflow.into()),
295            });
296        }
297
298        if let Some(language) = language {
299            let base = language.to_lowercase();
300            let base = base.split('-').next().unwrap_or("");
301            let chosen = if matches!(base, "en" | "eng" | "english") {
302                ModelName::English
303            } else {
304                ModelName::Multilingual
305            };
306            return Ok(RouteDecision {
307                model: chosen.as_str().into(),
308                repo: self.repo_string(chosen),
309                reason: format!("explicit lang={language:?}"),
310                detection: None,
311                workflow: workflow.map(Into::into),
312            });
313        }
314
315        let detection = lang::analyse(state);
316        let (chosen, reason) = if detection.script == "unknown" {
317            (
318                self.default,
319                format!(
320                    "no letters detected in state; using default ({})",
321                    self.default.as_str()
322                ),
323            )
324        } else if detection.script != "latin" {
325            (
326                ModelName::Multilingual,
327                format!(
328                    "non-Latin script ({}, {:.0}% of letters); the English checkpoint cannot read it",
329                    detection.script,
330                    100.0 * detection.non_latin_fraction
331                ),
332            )
333        } else if !detection.is_english {
334            (
335                ModelName::Multilingual,
336                format!(
337                    "Latin script but language looks like {:?}, not English",
338                    detection.language.clone().unwrap_or_default()
339                ),
340            )
341        } else {
342            (ModelName::English, "English Latin text".to_string())
343        };
344
345        Ok(RouteDecision {
346            model: chosen.as_str().into(),
347            repo: self.repo_string(chosen),
348            reason,
349            detection: Some(detection),
350            workflow: workflow.map(Into::into),
351        })
352    }
353
354    /// Route, then answer every question on the chosen checkpoint.
355    ///
356    /// The chosen checkpoint must already be attached; see [`Router::attach`],
357    /// or the `hub` feature for automatic downloads.
358    pub fn predict(
359        &mut self,
360        state: &Value,
361        questions: &Map<String, Value>,
362    ) -> Result<crate::agent::Prediction> {
363        let decision = self.route(state, Some(questions), None, None, None)?;
364        let chosen = ModelName::parse(&decision.model)?;
365        let agent = self.get(chosen).ok_or_else(|| {
366            Error::Config(format!(
367                "checkpoint {:?} is not loaded; attach it first",
368                decision.model
369            ))
370        })?;
371        self.evict();
372        let mut prediction = agent.predict_map(state, questions)?;
373        prediction.routing = Some(decision.to_value());
374        Ok(prediction)
375    }
376
377    fn evict(&mut self) {
378        while self.order.len() > self.max_loaded {
379            let victim = self.order.remove(0);
380            self.agents.remove(&victim);
381        }
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use serde_json::json;
389
390    fn questions(ids: &[&str]) -> Map<String, Value> {
391        ids.iter()
392            .map(|id| (id.to_string(), json!({"type": "noul", "instructions": "x"})))
393            .collect()
394    }
395
396    #[test]
397    fn aliases_resolve() {
398        assert_eq!(ModelName::parse("en").unwrap(), ModelName::English);
399        assert_eq!(ModelName::parse("ML").unwrap(), ModelName::Multilingual);
400        assert_eq!(
401            ModelName::parse("typed_decisions").unwrap(),
402            ModelName::TypedDecisions
403        );
404        assert!(ModelName::parse("nope").is_err());
405    }
406
407    #[test]
408    fn explicit_model_wins() {
409        let router = Router::new();
410        let decision = router
411            .route(&json!("hello"), None, Some("multilingual"), None, None)
412            .unwrap();
413        assert_eq!(decision.model, "multilingual");
414        assert!(decision.reason.contains("explicit model"));
415    }
416
417    #[test]
418    fn non_latin_routes_to_multilingual() {
419        let router = Router::new();
420        let decision = router
421            .route(&json!("Мой счёт был дважды списан"), None, None, None, None)
422            .unwrap();
423        assert_eq!(decision.model, "multilingual");
424        assert!(decision.reason.contains("non-Latin script"));
425    }
426
427    #[test]
428    fn english_routes_to_english() {
429        let router = Router::new();
430        let decision = router
431            .route(
432                &json!("I was charged twice and would like a refund please"),
433                None,
434                None,
435                None,
436                None,
437            )
438            .unwrap();
439        assert_eq!(decision.model, "english");
440    }
441
442    #[test]
443    fn workflow_detection_is_opt_in() {
444        let ids = questions(&["action", "category", "churn_risk", "needs_human", "urgency"]);
445        let off = Router::new();
446        assert_eq!(
447            off.route(&json!("hello"), Some(&ids), None, None, None)
448                .unwrap()
449                .model,
450            "english"
451        );
452
453        let on = Router::new().auto_task_detection(true);
454        let decision = on
455            .route(&json!("hello"), Some(&ids), None, None, None)
456            .unwrap();
457        assert_eq!(decision.model, "typed-decisions");
458        assert_eq!(decision.workflow.as_deref(), Some("customer_service"));
459    }
460
461    #[test]
462    fn partial_workflow_match_is_ignored() {
463        let ids = questions(&["urgency", "category"]);
464        assert_eq!(match_typed_decisions_workflow(&ids), None);
465    }
466
467    #[test]
468    fn explicit_language_routes_without_detection() {
469        let router = Router::new();
470        let decision = router
471            .route(&json!("hello"), None, None, None, Some("de-DE"))
472            .unwrap();
473        assert_eq!(decision.model, "multilingual");
474        assert!(decision.detection.is_none());
475    }
476
477    #[test]
478    fn unknown_script_uses_the_default() {
479        let router = Router::new().default_model(ModelName::Multilingual);
480        let decision = router
481            .route(&json!("12345"), None, None, None, None)
482            .unwrap();
483        assert_eq!(decision.model, "multilingual");
484        assert!(decision.reason.contains("default"));
485    }
486}
laya/router.rs

laya/
router.rs
