NeMo SwitchyardにTypeSafe(Jev)をclassifierとして組み込むライブラリを実装してみた

クラスメソッドマレーシアの森永です。

前回の記事「モデルルーティングをTypeSafe(Jev)に置き換えたら、どれくらい速く・安くなるか試してみた」で、TypeSafeのChoiceプリミティブを単体で叩いてみました。モデルルーティングが抱えていた、速いけど高い、安いけど遅いというトレードオフを、これで解消できそうな手応えがありました。ただ、実際にSwitchyardのlibsyに組み込むところまではやらずに終わっていました。
このまま終わるのは気持ち悪いので、Switchyardに実際にJevをclassifierとして組み込むライブラリを実装してみました。

Switchyardのリポジトリを覗いてみると、Issue #723 として、こちらが考えていたのとほぼ同じ内容のFeature Requestが立てられていました。

• 提案内容: TypeSafeのSystem OneをSwitchyardのルーティングに統合する
• Issue内に記載されたベンチマーク: TypeSafe 精度10/10・平均281ms・p95 375ms、GPT-5.6-sol 精度10/10・平均1653ms・p95 2576ms
• 実装要件: 正規化した会話状態をTypeSafeに型安全なChoiceとして送る。各ターゲットに安定したラベルと意味基準を定義する。信頼度が高い場合のみルーティングし、低い場合はフォールバックする
• セキュリティ要件: 認証情報は環境変数から取得し、TOML設定やログには残さない
• アーキテクチャ上の制約: 「switchyard-libsyはI/O-freeな設計を維持する。外部HTTP判定ステップ、またはランナー所有の分類プロバイダーとして実装し、libsyから直接HTTP呼び出しを行わない」

というわけで、今回はこのIssueの要件に沿う形で実装することにしました。

Switchyardのアーキテクチャを調べる

まずlibsy(ルーティングアルゴリズム本体)のコードを読んで、制約を具体的に把握しました。

libsyを読むと確かにI/O-freeな設計になっています。Algorithm::routeはdriver.call_model(request, models)を呼ぶだけで、実際のHTTP呼び出しはlibsy-llm-client側の実装が担当します。

もう一つの壁は、既存のlibsy-llm-clientのBackendが想定している通信形式(wire format)の種類です。

// crates/libsy-llm-client/src/backend.rs
enum Backend {
    OpenAiChat,
    OpenAiResponses,
    Anthropic,
}

TypeSafeのPOST /v1/systemoneは、このいずれの形式でもありません。チャット補完でもレスポンスAPIでもなく、{state, model, questions}という独自のリクエスト/レスポンス形式です。switchyard-translationにTypeSafeを4つ目のwire formatとして追加する案も検討しましたが、こちらはバッファリング・ストリーミング両対応のコーデックを新規に書く必要があり、TypeSafeの呼び出しがそもそもストリーミングされない(1回のリクエストで即座に判定が返る)ことを考えると相性が悪いように思いました。

そこで既存のllm-routing相当の実装、algorithms/llm_class.rsのLlmTaskClassifierを読みました。これはjudgeというcategoryのモデルに対して、JSON Schemaで構造化出力を強制したチャット補完リクエストを作り、driver.call_model経由で投げ、返ってきたテキストをJSONとしてパースしてポリシー判定する、という実装です。

実装方針

Issueが示す2つの選択肢(外部HTTP判定ステップと、ランナー所有の分類プロバイダー)を両方満たす形で、次のように実装しました。

libsy側: I/O-freeなポートと分類ロジック

まず、HTTPへの依存を一切持たないトレイトをlibsyに定義しました。これがIssueの言う外部HTTP判定ステップの入り口になります。

// crates/libsy/src/algorithms/util/typesafe_provider.rs
#[async_trait]
pub trait TypeSafeProvider: Send + Sync {
    async fn classify(
        &self,
        input: TypeSafeClassifierInput,
        options: &[TypeSafeOption],
    ) -> Result<TypeSafeVerdict, TypeSafeProviderError>;
}

libsyはこのトレイトに依存するだけで、実際にHTTPを叩く実装は一切持ちません。TypeSafeOption(label + description)、TypeSafeClassifierInput(question + context)、TypeSafeVerdict(label + confidence)は、いずれもプレーンなデータ型です。

その上で、LlmTaskClassifierと同じ骨格を持つTypeSafeTaskClassifierを実装しました。

// crates/libsy/src/algorithms/type_safe_class.rs
pub struct TypeSafeClassifierConfig {
    pub options: Vec<TypeSafeOption>,
    pub question: String,
    pub base_threshold: f64,
    pub default_target: Category,
    pub classify_trigger: ClassifyTrigger,
    pub message_hash_fallback: bool,
    pub recent_turn_window: Option<usize>,
}

provider(TypeSafeProvider)を呼び、confidenceがbase_threshold未満だったり、providerがエラーを返したり、返ってきたラベルがどのCategoryにも解決できなかったりした場合は、すべて判定不能とみなしてDefaultCategoryClassifier側に流します。何が起きてもErrを返さない、fail-open設計です。

会話のウィンドウ処理(trim_messages/task_messages)やAffinityRouterの構築ロジックは、既存のllm_class.rsのものをそのままpub(crate)にして再利用し、重複した実装を持たないようにしました。

switchyard-typesafe-client: 実際にHTTPを叩く新規クレート

Issueの言うランナー所有の分類プロバイダーが、この新規クレートです。libsy-llm-clientの兄弟クレートとして、switchyard-typesafe-clientを追加しました。

// crates/switchyard-typesafe-client/src/lib.rs
#[async_trait]
impl TypeSafeProvider for TypeSafeHttpClient {
    async fn classify(
        &self,
        input: TypeSafeClassifierInput,
        options: &[TypeSafeOption],
    ) -> Result<TypeSafeVerdict, TypeSafeProviderError> {
        // { state, model, questions: { "route": { type: "choice", instructions, criteria } } }
        // を POST {base_url}/v1/systemone に送り、
        // answers.route.{choice, confidence} を TypeSafeVerdict に詰め替える
        ...
    }
}

リクエスト/レスポンスの形式は前回記事の検証で確認したものと同一です。APIキーはTypeSafeHttpClient::from_env(変数名)で環境変数からのみ取得する設計にし、TOML設定ファイルには一切書けないようにしました。

switchyard-runner: TOML設定への配線

最後に、実際のデプロイ設定(TOML)からこの機能を使えるように、switchyard-runnerに配線しました。

[type_safe_client]
api_key_env = "TYPESAFE_API_KEY"

[routes.switchyard]
type = "type_safe_classifier"
default_target = "efficient"
base_threshold = 0.6
question = "Which model tier does this conversation need?"

[routes.switchyard.options]
capable = "Needs multi-step reasoning, ambiguous instructions, or high-stakes correctness."
efficient = "A short, well-specified request."

[routes.switchyard.models]
capable = ["gpt-5.5"]
efficient = ["gpt-5.5-mini"]
any = ["gpt-5.5", "gpt-5.5-mini"]

[type_safe_client]テーブルが無いのにtype = "type_safe_classifier"のルートを定義すると、設定エラーとして起動時に弾かれます。optionsの各キーはmodelsにも同名のグループが必要、といったバリデーションも、既存のllm_classifierの設定検証ロジックを共通化して実装しました。

ベンチマーク

前回記事のベンチマークは、TypeSafeのAPIをPython(urllib)から直接叩いた場合の実測でした。今回は、実際にSwitchyardへ組み込んで、同じベンチマークをもう一度実施しました。ライブラリ化・Rust実装への置き換えによって速度が落ちていないかを確認するのが目的です。

検証方法

前回と条件は完全に同一です。NeMo Switchyardのllm-routingを模した4パターン(simple/medium/complex/reasoningそれぞれを期待する会話サマリ)を用意し、各10回、合計40回、TypeSafeHttpClient::classifyを呼び出しました。

結果

40回すべてが期待したティアと一致しました(正解率40/40)。

※ simpleのp95(0.747秒)は10回中1回だけの外れ値です。実行したプロセスの最初の呼び出しにあたり、TLSハンドシェイクやコネクション確立のコストが乗ったと考えられます。それ以外の9回は0.2〜0.4秒台に収まっており、median/meanの数値にはほとんど影響していません。

前回記事の直接API実測、およびIssue #723に記載されたTypeSafe単体のベンチマーク(平均281ms・p95 375ms)と並べてみます。

前回記事の実測(0.64〜0.69秒)より、今回の実測(0.25〜0.32秒)の方がかなり速く出ています。これは実装の優劣ではなく、単純にクライアント実行環境(Rust + reqwestか、Python + urllibか)やその時点のネットワーク・API側の状況の違いによるものと考えられ、両者を厳密に比較するための条件を揃えたベンチマークではありません。ただ、Issue #723記載のTypeSafe単体ベンチマーク(平均281ms)とは近い値になっており、ライブラリ化・Rust実装への置き換えによって、TypeSafe本来の速さを損なわずに使えていることは確認できたと考えています。(今回コストは計算していませんが、今回のベンチマークで使ったstateテキストは前回記事と同一なので、同じコストと考えられます。)

まとめ

実際のTypeSafe APIとの疎通確認、そして前回記事と同条件のベンチマークを実施し、Jev本来の速さ(median 0.25〜0.28秒、40/40の分類精度)が損なわれていないことを確認しました。
NeMo SwitchyardのようなLLMルーティング基盤を運用する際は、Jevがかなり有力な選択肢になりそうな手応えを感じました。

PR 出したところいろんな議論が始まって楽しいです。