Skip to main content

switchyard_protocol/
category.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Category is a group of models
5
6use std::{collections::HashMap, str::FromStr};
7
8use crate::ModelId;
9
10/// A group of models
11#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
12pub enum Category {
13    /// When the category doesn't matter: Random, Passthrough, etc.
14    Any,
15    /// High accuracy and cost models.
16    Capable,
17    /// Lower accuracy and cost models.
18    Efficient,
19    /// Models the algorithm can use to decide.
20    Judge,
21}
22
23impl Category {
24    /// Convenience function to create a HashMap suitable for passing to `run_stream` and family.
25    pub fn to_map(category: Category, names: &[&str]) -> HashMap<Category, Vec<ModelId>> {
26        [(
27            category,
28            names.iter().map(|name| ModelId::from(*name)).collect(),
29        )]
30        .into()
31    }
32}
33
34impl FromStr for Category {
35    type Err = String;
36    fn from_str(s: &str) -> Result<Self, Self::Err> {
37        let c = match s {
38            "capable" => Self::Capable,
39            "efficient" => Self::Efficient,
40            "judge" => Self::Judge,
41            "any" => Self::Any,
42            x => {
43                return Err(format!("Invalid Category '{x}'"));
44            }
45        };
46        Ok(c)
47    }
48}