feat(llm-classifier): custom schema support + multi target (>2) routing support - #275
Conversation
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
Signed-off-by: ayushag <ayushag@nvidia.com>
|
WalkthroughSummaryThe PR adds reusable structured LLM judges, JSON-schema verdict validation, custom JSON Pointer target selection, and explicit capability, escalation, and custom classifier modes. Server configuration, tests, exports, and documentation now support these modes. ChangesLLM classifier modes
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (6)
crates/switchyard-server/tests/server.rs (1)
817-861: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the classifier call count and the untouched targets.
The loop asserts the selected model for each turn. It does not assert that the classifier target was called once per turn, so a cached or skipped judge call would still pass. It also does not assert that
model/weakandmodel/middleare never selected, which is the core claim of routing over more than two targets.💚 Proposed additional assertions
let calls = upstream.calls.lock().await; + let judge_calls = calls + .iter() + .filter(|call| call["model"] == "model/classifier") + .count(); + assert_eq!(judge_calls, 2, "{calls:?}"); + assert!( + !calls + .iter() + .any(|call| call["model"] == "model/weak" || call["model"] == "model/middle"), + "{calls:?}" + ); let judge_call = calls .iter() .find(|call| call["model"] == "model/classifier") .ok_or("custom classifier target was not called")?;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-server/tests/server.rs` around lines 817 - 861, Extend the test assertions after the request loop to verify the classifier target model/classifier call count equals the number of turns, preventing cached or skipped classifications from passing. Also assert that the selected-model results never include model/weak or model/middle, while preserving the existing assertions for model/premium and model/strong and the classifier prompt/schema checks.crates/libsy/src/algorithms/util/target_selector.rs (1)
107-112: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a rejection test for a root selector.
newrejects a root pointer at lines 31-35, but no test covers it. An empty selector ("") is a valid JSON Pointer that resolves to the whole verdict, so this branch is the only guard against it.💚 Proposed additional test
#[test] fn an_invalid_json_pointer_is_rejected() { let result = TargetSelectorPolicy::new("/target~2name", BTreeMap::new()); assert!(matches!(result, Err(LibsyError::AlgorithmError { message }) if message.contains("valid JSON Pointer"))); } + + #[test] + fn a_root_selector_is_rejected() { + let result = TargetSelectorPolicy::new("", BTreeMap::new()); + assert!(matches!(result, Err(LibsyError::AlgorithmError { message }) + if message.contains("must identify a response field"))); + } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/util/target_selector.rs` around lines 107 - 112, Add a test alongside an_invalid_json_pointer_is_rejected that calls TargetSelectorPolicy::new with an empty selector and asserts it returns the expected LibsyError::AlgorithmError for rejecting a root selector. Keep the test focused on verifying that the valid JSON Pointer "" is not accepted because it would select the entire verdict.crates/libsy/src/algorithms/llm_class.rs (1)
605-708: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd test coverage for
new_customvalidation in this module.
new_customadds several new failure paths: fewer than two targets, an untrimmed or empty label, a duplicated label, a duplicated resolvedsemantic_name, and adefault_targetthat is not a configured label. The tests in this file cover only the capability and escalation constructors, and the server integration test covers only the success and fallback paths. Add unit tests for these five rejections so a future change cannot relax them silently.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/llm_class.rs` around lines 605 - 708, Add unit tests in the llm_class test module covering each new_custom rejection: fewer than two targets, empty or surrounding-whitespace labels, duplicate labels, duplicate target semantic_name values, and a default_target absent from the configured labels. Reuse the existing test target/configuration helpers where available, invoke new_custom for each case, and assert that construction returns an error while preserving the existing successful and fallback coverage.crates/switchyard-server/src/config.rs (2)
946-957: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend the mode tests to cover inference and the remaining rejections.
The new test covers one rejection. The mode logic added in
classifier_modehas several other branches that no test exercises:
- An omitted
modewithescalationpresent selects escalation, and an omittedmodewithoutescalationselects capability.mode = "capability"withescalationpresent is rejected.mode = "escalation"withbase_thresholdpresent is rejected.- Any non-custom mode with
targets,default_target,response_schema, orpolicyis rejected.mode = "capability"withoutbase_thresholdreports "requires base_threshold", which replaces the previous serde "missing field" error.Add cases for these so the mode contract is pinned.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-server/src/config.rs` around lines 946 - 957, Extend the mode-related tests around mode_custom_rejects_capability_fields to cover classifier_mode inference and every listed rejection: omitted mode with and without escalation, capability with escalation, escalation with base_threshold, non-custom modes using custom-only fields, and capability without base_threshold expecting the classifier’s “requires base_threshold” message. Reuse existing configuration fixtures and error_message helpers, adding focused cases that assert the documented selected mode or error text.
390-394: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the mode inference and the
mode.is_some()gate.Two decisions in
classifier_modeare not obvious from the code. First, lines 390-394 infer the mode from the presence ofescalationso that configurations written beforemodeexisted keep working. Second, line 449 gates the capability-settings rejection onmode.is_some(), so a legacy escalation route that still carriesbase_thresholdis accepted and those settings are ignored instead of rejected. Add short comments stating both intents, so a later change does not remove the compatibility path or tighten it accidentally.As per coding guidelines: "add concise comments for module intent, private helpers with non-obvious behavior, important tests, and complex validation, routing, configuration, async, lifecycle, or concurrency logic".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-server/src/config.rs` around lines 390 - 394, Add concise comments around the selected_mode inference in classifier_mode explaining that escalation presence preserves compatibility for configurations created before mode existed. Also comment the mode.is_some() gate on capability-settings rejection, noting that legacy escalation routes may retain base_threshold and those settings are intentionally ignored rather than rejected.Source: Coding guidelines
crates/libsy/src/algorithms/util/classifier_contract.rs (1)
64-85: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winValidate strict JSON Schema constraints at construction.
from_inner_schemaalways sends"strict": true, butcompile_schemachecks only generic JSON Schema validity. OpenAI-style strict providers require every object to setadditionalProperties: falseand list every property inrequired. Invalid strict schemas can pass construction, cause every judge request to fail, and route every turn todefault_targetwhen the failure becomesNone.Validate these constraints recursively in
from_inner_schema, or make strict mode configurable and document the required schema shape.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/util/classifier_contract.rs` around lines 64 - 85, Update from_inner_schema to validate the complete schema recursively against the strict JSON Schema constraints required by the generated "strict": true response format: every object must set additionalProperties to false and include all declared property names in required, including nested objects and objects inside arrays or combinators. Reject violations with the existing Result error path before calling from_response_format, while preserving compilation and construction for valid schemas.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/reference/toml_schema.md`:
- Line 134: Update the policy table entry in toml_schema.md to explicitly
identify selector as the field that accepts the JSON Pointer, while
distinguishing it from the policy type value target_selector. Keep the existing
example and table structure unchanged.
In `@docs/routing_algorithms/llm_classifier_routing.md`:
- Around line 110-122: Update the custom multi-target routing example around the
routes.smart configuration to define the referenced fast, balanced, reasoning,
and premium targets with valid target tables, or explicitly indicate that this
block replaces the earlier target definitions. Ensure the copied configuration
passes target-reference validation.
---
Nitpick comments:
In `@crates/libsy/src/algorithms/llm_class.rs`:
- Around line 605-708: Add unit tests in the llm_class test module covering each
new_custom rejection: fewer than two targets, empty or surrounding-whitespace
labels, duplicate labels, duplicate target semantic_name values, and a
default_target absent from the configured labels. Reuse the existing test
target/configuration helpers where available, invoke new_custom for each case,
and assert that construction returns an error while preserving the existing
successful and fallback coverage.
In `@crates/libsy/src/algorithms/util/classifier_contract.rs`:
- Around line 64-85: Update from_inner_schema to validate the complete schema
recursively against the strict JSON Schema constraints required by the generated
"strict": true response format: every object must set additionalProperties to
false and include all declared property names in required, including nested
objects and objects inside arrays or combinators. Reject violations with the
existing Result error path before calling from_response_format, while preserving
compilation and construction for valid schemas.
In `@crates/libsy/src/algorithms/util/target_selector.rs`:
- Around line 107-112: Add a test alongside an_invalid_json_pointer_is_rejected
that calls TargetSelectorPolicy::new with an empty selector and asserts it
returns the expected LibsyError::AlgorithmError for rejecting a root selector.
Keep the test focused on verifying that the valid JSON Pointer "" is not
accepted because it would select the entire verdict.
In `@crates/switchyard-server/src/config.rs`:
- Around line 946-957: Extend the mode-related tests around
mode_custom_rejects_capability_fields to cover classifier_mode inference and
every listed rejection: omitted mode with and without escalation, capability
with escalation, escalation with base_threshold, non-custom modes using
custom-only fields, and capability without base_threshold expecting the
classifier’s “requires base_threshold” message. Reuse existing configuration
fixtures and error_message helpers, adding focused cases that assert the
documented selected mode or error text.
- Around line 390-394: Add concise comments around the selected_mode inference
in classifier_mode explaining that escalation presence preserves compatibility
for configurations created before mode existed. Also comment the mode.is_some()
gate on capability-settings rejection, noting that legacy escalation routes may
retain base_threshold and those settings are intentionally ignored rather than
rejected.
In `@crates/switchyard-server/tests/server.rs`:
- Around line 817-861: Extend the test assertions after the request loop to
verify the classifier target model/classifier call count equals the number of
turns, preventing cached or skipped classifications from passing. Also assert
that the selected-model results never include model/weak or model/middle, while
preserving the existing assertions for model/premium and model/strong and the
classifier prompt/schema checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 88bd5c89-9361-4215-a0f1-f622210021ab
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (16)
Cargo.tomlcrates/libsy/Cargo.tomlcrates/libsy/src/algorithms/llm_class.rscrates/libsy/src/algorithms/util.rscrates/libsy/src/algorithms/util/classifier_contract.rscrates/libsy/src/algorithms/util/escalation.rscrates/libsy/src/algorithms/util/llm_judge.rscrates/libsy/src/algorithms/util/target_selector.rscrates/libsy/src/lib.rscrates/switchyard-server/README.mdcrates/switchyard-server/src/config.rscrates/switchyard-server/tests/server.rsdocs/getting_started.mddocs/reference/toml_schema.mddocs/routing_algorithms/escalation_router_routing.mddocs/routing_algorithms/llm_classifier_routing.md
Signed-off-by: ayushag <ayushag@nvidia.com>
|
Can we add an example 2+ model prompt in examples section? |
This can be a follow-up PR |
Signed-off-by: ayushag <ayushag@nvidia.com>
3 participants
Why
Classifier prompts are configurable, but response schemas and verdict-to-target mapping were still hard-coded. Every new verdict shape required another Rust judge implementation, and the server route could only choose between two named tiers.
What
capability,escalation, andcustomclassifier modes.target_selectorpolicy for routing from a verdict field.default_targetwhen the judge fails, returns invalid JSON, fails schema validation, or selects an unknown label.modeis omitted.How
The classifier path is split into four small responsibilities:
ClassifierInputchooses what the judge sees.TaskInputhandles capability and custom routing;EscalationInputbuilds the condensed completed trajectory.ClassifierContractrenders the system prompt, inserts the inner schema into the provider response-format wrapper, and owns the compiled validator for custom verdicts.StructuredJudge<I, D>assembles the judge request once.SerdeDecoder<V>parses the fixed capability and escalation verdict structs, whileJsonSchemaDecoderparses a dynamic JSON value and validates it against the configured schema.JudgePolicyconverts the decoded verdict into a routing decision. Capability and escalation keep their typed policies; custom mode appliesTargetSelectorPolicyto a configured selector.This keeps request construction and model calling shared without forcing every classifier into a dynamic JSON representation. Fixed contracts stay typed. Custom contracts pay for runtime schema validation only when that mode is configured.
The server preserves the flat TOML interface, validates fields against the selected mode, and normalizes the route into typed capability, escalation, or custom construction settings. The existing
FallThroughcomposition still owns affinity and default fallback behavior.Where to Start Review
crates/libsy/src/algorithms/util/llm_judge.rsforStructuredJudge, input shaping, and decoder boundaries.crates/libsy/src/algorithms/util/classifier_contract.rsfor schema wrapping, compilation, and verdict validation.crates/libsy/src/algorithms/llm_class.rsfor mode construction, policy wiring, and fallback composition.crates/libsy/src/algorithms/util/target_selector.rsfor deterministic multi-target selection.crates/switchyard-server/src/config.rsfor mode selection and TOML validation.crates/switchyard-server/tests/server.rsfor request-level compatibility and four-target routing coverage.Test Plan
cargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warningscargo test --workspacemodecd docs && make publishCloses
Closes SWITCH-1180