Skip to content

feat(libsy): pin the classifier target across a user turn #487

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ryan-lempka wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(libsy): pin the classifier target across a user turn #487

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion benchmark/server-configs/tb-lite-llm-classifier-opus-kimi-gemini.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,5 +31,5 @@ strong_target = "strong"
weak_target = "weak"
base_threshold = 0.5
threshold_step = 0.0
session_affinity = true
classify_trigger = "new_session"
message_hash_fallback = true
8 changes: 4 additions & 4 deletions crates/libsy-llm-client/tests/observability.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ use tracing_subscriber::layer::{Context as LayerContext, SubscriberExt};
use tracing_subscriber::registry::LookupSpan;

use switchyard_libsy::{
AffinityRouter, Algorithm, Classifier, Driver, LibsyError, LlmClassifierConfig,
LlmTaskClassifier, PickerMode, RoutingOutcome, StageRouter, StageRouterConfig, Step,
TaskClassifierConfig,
AffinityRouter, Algorithm, Classifier, ClassifyTrigger, Driver, LibsyError,
LlmClassifierConfig, LlmTaskClassifier, PickerMode, RoutingOutcome, StageRouter,
StageRouterConfig, Step, TaskClassifierConfig,
};
use switchyard_llm_client::{ClientRouter, RunObservation, RunObserver};
use switchyard_protocol::ModelId;
Expand Down Expand Up @@ -676,7 +676,7 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback()
capable_target: "affinity-fallback-strong".into(),
config: TaskClassifierConfig {
base_threshold: 0.5,
session_affinity: true,
classify_trigger: ClassifyTrigger::NewSession,
..TaskClassifierConfig::default()
},
})?) as Arc<dyn Algorithm>;
Expand Down
66 changes: 35 additions & 31 deletions crates/libsy/src/algorithms/llm_class.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ use serde_json::Value;
use switchyard_protocol::{ContentBlock, Message, ModelId, Role};

use super::fall_through::{DefaultTarget, FallThrough};
use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
use super::util::affinity::AffinityRouter;
use super::util::classifier_contract::{
ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
Expand All @@ -23,6 +22,8 @@ use super::util::llm_judge::{
SerdeDecoder, StructuredJudge,
};
use super::util::target_selector::TargetSelectorPolicy;
use super::util::turn_pin::{ClassifyTrigger, TurnPin};
use super::util::{DEFAULT_JUDGE_MAX_OUTPUT_TOKENS, decisive};
use crate::core::algorithm::{self, Algorithm, Driver};
use crate::core::classifier::{Classification, Classifier, Score};
use crate::core::state::{State, StateValue};
Expand Down Expand Up @@ -231,8 +232,8 @@ pub struct TaskClassifierConfig {
/// Supported verdicts use `base_threshold`, uncertain and unmatched verdicts use one
/// step, and unsupported verdicts use two steps.
pub threshold_step: f64,
/// Enables session affinity before the judge-backed classifier.
pub session_affinity: bool,
/// How often the classifier re-decides this session's target.
pub classify_trigger: ClassifyTrigger,
/// Uses the first user message as the SessionKey for sticky routing when session metadata is unavailable.
pub message_hash_fallback: bool,
/// Trailing conversation turns the judge sees on top of the client
Expand All @@ -256,7 +257,7 @@ struct TaskClassifierConfigWire {
#[serde(default)]
threshold_step: f64,
#[serde(default)]
session_affinity: bool,
classify_trigger: ClassifyTrigger,
#[serde(default)]
message_hash_fallback: bool,
#[serde(default)]
Expand All @@ -283,7 +284,7 @@ impl<'de> Deserialize<'de> for TaskClassifierConfig {
Ok(Self {
base_threshold: wire.base_threshold,
threshold_step: wire.threshold_step,
session_affinity: wire.session_affinity,
classify_trigger: wire.classify_trigger,
message_hash_fallback: wire.message_hash_fallback,
recent_turn_window: wire.recent_turn_window,
contract,
Expand All @@ -301,7 +302,7 @@ impl Default for TaskClassifierConfig {
Self {
base_threshold: 0.0,
threshold_step: 0.0,
session_affinity: false,
classify_trigger: ClassifyTrigger::default(),
message_hash_fallback: false,
recent_turn_window: None,
contract: ClassifierContractConfig::default(),
Expand Down Expand Up @@ -342,9 +343,10 @@ impl TaskClassifierConfig {
message: "max_output_tokens must be at least 1".to_string(),
});
}
if self.message_hash_fallback && !self.session_affinity {
if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
return Err(LibsyError::AlgorithmError {
message: "message_hash_fallback requires session_affinity".to_string(),
message: "message_hash_fallback requires classify_trigger = new_session"
.to_string(),
});
}
Ok(())
Expand Down Expand Up @@ -379,8 +381,8 @@ pub struct CustomClassifierConfig {
pub response_schema: Value,
/// Deterministic policy applied after the verdict passes schema validation.
pub policy: CustomClassifierPolicy,
/// Enables session affinity before the judge-backed classifier.
pub session_affinity: bool,
/// How often the classifier re-decides this session's target.
pub classify_trigger: ClassifyTrigger,
/// Uses the first user message when session metadata is unavailable.
pub message_hash_fallback: bool,
/// Trailing conversation turns shown to the classifier judge.
Expand All @@ -400,7 +402,7 @@ impl CustomClassifierConfig {
prompt: prompt.into(),
response_schema,
policy,
session_affinity: false,
classify_trigger: ClassifyTrigger::default(),
message_hash_fallback: false,
recent_turn_window: None,
max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
Expand All @@ -413,9 +415,10 @@ impl CustomClassifierConfig {
message: "max_output_tokens must be at least 1".to_string(),
});
}
if self.message_hash_fallback && !self.session_affinity {
if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
return Err(LibsyError::AlgorithmError {
message: "message_hash_fallback requires session_affinity".to_string(),
message: "message_hash_fallback requires classify_trigger = new_session"
.to_string(),
});
}
Ok(())
Expand Down Expand Up @@ -454,13 +457,6 @@ fn streak(state: &State) -> u32 {
}
}

fn decisive(target: &ModelId) -> Classification {
Classification::Scores(vec![Score {
target: target.clone(),
confidence: 1.0,
}])
}

fn assistant_message(response: &AggLlmResponse) -> Message {
Message {
role: Role::Assistant,
Expand Down Expand Up @@ -592,7 +588,7 @@ pub struct LlmTaskClassifier {

struct ClassifierRouteConfig {
default_target: ModelId,
session_affinity: bool,
classify_trigger: ClassifyTrigger,
message_hash_fallback: bool,
}

Expand Down Expand Up @@ -687,7 +683,7 @@ impl LlmTaskClassifier {
config.validate()?;
let contract = Self::load_capability_contract(&config.contract)?;
let targets = vec![efficient_target.clone(), capable_target.clone()];
let session_affinity = config.session_affinity;
let classify_trigger = config.classify_trigger;
let message_hash_fallback = config.message_hash_fallback;
let classifier = Arc::new(TaskClassifier {
classifier: JudgeClassifier::new(
Expand Down Expand Up @@ -715,7 +711,7 @@ impl LlmTaskClassifier {
inner,
ClassifierRouteConfig {
default_target: classifier.capable_target.clone(),
session_affinity,
classify_trigger,
message_hash_fallback,
},
)
Expand Down Expand Up @@ -772,7 +768,7 @@ impl LlmTaskClassifier {
prompt,
response_schema,
policy,
session_affinity,
classify_trigger,
message_hash_fallback,
recent_turn_window,
max_output_tokens,
Expand Down Expand Up @@ -801,7 +797,7 @@ impl LlmTaskClassifier {
classifier,
ClassifierRouteConfig {
default_target: default_name,
session_affinity,
classify_trigger,
message_hash_fallback,
},
)
Expand Down Expand Up @@ -853,16 +849,24 @@ impl LlmTaskClassifier {
config: ClassifierRouteConfig,
) -> Result<Self> {
algorithm::ensure_model_is_target(&targets, &config.default_target)?;
if config.message_hash_fallback && !config.session_affinity {
if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession {
return Err(LibsyError::AlgorithmError {
message: "message_hash_fallback requires session_affinity".to_string(),
message: "message_hash_fallback requires classify_trigger = new_session"
.to_string(),
});
}
// Wraps the classifier rather than the route, so the pin also holds when this is
// embedded in another cascade and only `score` is called.
let inner = if config.classify_trigger == ClassifyTrigger::UserTurn {
Arc::new(TurnPin::new(inner)) as Arc<dyn Classifier<State>>
} else {
inner
};
// Affinity comes first so a retained assignment short-circuits the judge call.
// Note: when this classifier is embedded inside another cascade (e.g. StageRouter)
// the affinity processor never fires — only the inner score() is called.
let mut route = FallThrough::<State>::new_with_state(targets).with_name(ALGORITHM_NAME);
if config.session_affinity {
if config.classify_trigger == ClassifyTrigger::NewSession {
let affinity = if config.message_hash_fallback {
AffinityRouter::new().with_message_hash_fallback()
} else {
Expand Down Expand Up @@ -1206,14 +1210,14 @@ mod tests {
}

#[tokio::test]
async fn classifier_config_enables_session_affinity() -> Result<()> {
async fn classifier_config_enables_new_session_trigger() -> Result<()> {
let recorder = Arc::new(Recorder::default());
let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
judge_target: ModelId::from("judge"),
efficient_target: ModelId::from("efficient"),
capable_target: ModelId::from("capable"),
config: TaskClassifierConfig {
session_affinity: true,
classify_trigger: ClassifyTrigger::NewSession,
..test_config(TEST_THRESHOLD)
},
})?);
Expand All @@ -1234,7 +1238,7 @@ mod tests {
efficient_target: ModelId::from("efficient"),
capable_target: ModelId::from("capable"),
config: TaskClassifierConfig {
session_affinity: true,
classify_trigger: ClassifyTrigger::NewSession,
message_hash_fallback: true,
recent_turn_window: None,
..test_config(TEST_THRESHOLD)
Expand Down
2 changes: 1 addition & 1 deletion crates/libsy/src/algorithms/stage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ pub struct LlmFallback {
pub judge_target: ModelId,
/// Judge configuration. `recent_turn_window` is worth setting to this router's
/// `recent_window` so the judge reads the same span the signal scorer scored.
/// Note: `session_affinity` and `message_hash_fallback` have no effect here —
/// Note: `classify_trigger = new_session` and `message_hash_fallback` have no effect here —
/// the judge runs as a cascade classifier, not a standalone algorithm.
pub config: TaskClassifierConfig,
}
Expand Down
13 changes: 13 additions & 0 deletions crates/libsy/src/algorithms/util.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,19 @@ pub(crate) mod stage;
pub mod subagent;
pub(crate) mod target_selector;
pub(crate) mod tool_signals;
pub mod turn_pin;

use switchyard_protocol::ModelId;

use crate::core::classifier::{Classification, Score};

/// A single full-confidence recommendation.
pub(crate) fn decisive(target: &ModelId) -> Classification {
Classification::Scores(vec![Score {
target: target.clone(),
confidence: 1.0,
}])
}

/// Default completion budget for internal classifier and escalation judge calls.
pub(crate) const DEFAULT_JUDGE_MAX_OUTPUT_TOKENS: u64 = 4_096;
Loading
Loading