diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 970d8de85755..1fa98f13c255 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -426,13 +426,13 @@ class ComplexityRouterConfig(BaseModel): # Session affinity: pin the first turn's routed model for the rest of the session session_affinity: bool = Field( - default=True, + default=False, description=( "When True and a session_id is resolvable on the request, pin the model chosen on the " "session's first turn and reuse it for every later turn, skipping re-classification. " - "On by default so multi-turn sessions stay on one model, preserving provider prompt " - "caches and avoiding cross-model conversation-history errors. Set False to reclassify " - "every turn." + "Off by default so every turn is classified on its own merits and routed to the cheapest " + "adequate tier. Set True to keep a multi-turn session on one model, which preserves " + "provider prompt caches and avoids cross-model conversation-history errors." ), ) session_affinity_ttl_seconds: int = Field( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index cc73273450aa..3a94b1e0f85c 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2903,7 +2903,7 @@ async def test_complexity_scorer_logs_its_cause(self, mock_router_instance, basi class TestSessionAffinity: - """Test the session_affinity sticky-routing behavior (on by default).""" + """Test the session_affinity sticky-routing behavior (off by default).""" REASONING_MESSAGE = [ { @@ -2917,18 +2917,14 @@ class TestSessionAffinity: def session_affinity_config(self, basic_config) -> Dict: return {**basic_config, "session_affinity": True} - @pytest.fixture - def session_affinity_disabled_config(self, basic_config) -> Dict: - return {**basic_config, "session_affinity": False} - @staticmethod def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio - async def test_enabled_by_default_pins_model(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to True, so a shared session_id pins the - first turn's model and later turns reuse it instead of reclassifying.""" + async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to False, so a shared session_id must NOT + pin the first turn's model; every turn is classified on its own merits.""" assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( @@ -2944,19 +2940,17 @@ async def test_enabled_by_default_pins_model(self, mock_router_instance, basic_c model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" - assert second.model == "o1-preview" + assert second.model == "gpt-4o-mini" @pytest.mark.asyncio - async def test_can_be_disabled_reclassifies_every_turn( - self, mock_router_instance, session_affinity_disabled_config - ): - """Regression: session_affinity=False must still reclassify every turn even when a - shared session_id is present, so the opt-out keeps working.""" + async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): + """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the + first turn's model instead of reclassifying.""" mock_router_instance.cache = DualCache() router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config=session_affinity_disabled_config, + complexity_router_config=session_affinity_config, ) request_kwargs = self._request_kwargs("session-1") first = await router.async_pre_routing_hook( @@ -2966,7 +2960,7 @@ async def test_can_be_disabled_reclassifies_every_turn( model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" - assert second.model == "gpt-4o-mini" + assert second.model == "o1-preview" @pytest.mark.asyncio async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config): diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 0503c0c9c6d5..b83808b0728f 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -14,6 +14,7 @@ export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000; export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3; export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200; +export const DEFAULT_SESSION_AFFINITY = false; export interface ComplexityTiers { SIMPLE: string[]; @@ -45,6 +46,7 @@ export interface ComplexityRouterConfigValue { classifier_context_window_size?: number; classifier_context_per_turn_chars?: number; classifier_context_include_assistant_turns?: boolean; + session_affinity?: boolean; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; @@ -224,6 +226,32 @@ const ComplexityRouterConfig: React.FC = ({ ), children: , }, + { + key: "session-affinity", + label: ( + + Advanced: Session Affinity + + ), + children: ( + <> +
+ onChange({ ...value, session_affinity: sessionAffinity })} + aria-label="Pin a session to its first model" + /> + Pin a session to its first model +
+ + Off by default: every turn is classified on its own merits and routed to the cheapest adequate tier. + Turn this on to reuse the model chosen on a session's first turn for every later turn, which + preserves provider prompt caches and avoids cross-model conversation-history errors, at the cost of + keeping the whole session on the first turn's tier. + + + ), + }, { key: "response", label: ( diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index fa31c7d9c9df..f7cc9a1deaee 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -111,4 +111,40 @@ describe("AddAutoRouterTab", () => { expect(await screen.findByText("Please select a team to continue")).toBeInTheDocument(); expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); }); + + it("defaults a new router to session affinity off, matching the backend field default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + await user.click(screen.getByText("Advanced: Session Affinity")); + expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + session_affinity: false, + }); + }); + + it("carries session affinity turned on through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + await user.click(screen.getByText("Advanced: Session Affinity")); + await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + session_affinity: true, + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 4593f6a6a2ed..ea75bd8e2837 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -10,6 +10,7 @@ import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_m import ComplexityRouterConfig, { ComplexityRouterConfigValue, DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_SESSION_AFFINITY, DEFAULT_TIER_DISTANCE_PENALTY, } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; @@ -102,6 +103,7 @@ const AddAutoRouterTab: React.FC = ({ classifier_context_window_size: classifierContextWindowSize, classifier_context_per_turn_chars: classifierContextPerTurnChars, classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns, + session_affinity: sessionAffinity = DEFAULT_SESSION_AFFINITY, adaptive = false, adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, @@ -148,6 +150,7 @@ const AddAutoRouterTab: React.FC = ({ classifierContextWindowSize, classifierContextPerTurnChars, classifierContextIncludeAssistantTurns, + sessionAffinity, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index e939ce129047..9d784b579034 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -19,6 +19,7 @@ const baseParams: BuildComplexityRouterConfigParams = { classifierContextWindowSize: undefined, classifierContextPerTurnChars: undefined, classifierContextIncludeAssistantTurns: undefined, + sessionAffinity: false, customTechnicalKeywords: [], keywordTierRules: [], semanticMatchingEnabled: false, @@ -35,7 +36,12 @@ const baseParams: BuildComplexityRouterConfigParams = { describe("buildComplexityRouterConfig", () => { it("emits tiers, classifier_type, and escalation_keywords when nothing else is configured", () => { const config = buildComplexityRouterConfig(baseParams); - expect(config).toEqual({ tiers, classifier_type: "heuristic", escalation_keywords: ["LITELLM ESCALATE"] }); + expect(config).toEqual({ + tiers, + classifier_type: "heuristic", + session_affinity: false, + escalation_keywords: ["LITELLM ESCALATE"], + }); }); it("trims escalation keywords and drops blank entries", () => { @@ -219,6 +225,16 @@ describe("buildComplexityRouterConfig", () => { expect(config.return_raw_model_name).toBeUndefined(); }); + it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => { + const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true }); + expect(config.session_affinity).toBe(true); + }); + + it("writes session_affinity explicitly when off, so the stored config never relies on the backend default", () => { + const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: false }); + expect(config.session_affinity).toBe(false); + }); + it("includes return_raw_model_name when enabled", () => { const config = buildComplexityRouterConfig({ ...baseParams, returnRawModelName: true }); expect(config.return_raw_model_name).toBe(true); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 192e71b4597a..cd6c697b3776 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -15,6 +15,7 @@ export interface BuildComplexityRouterConfigParams { classifierContextWindowSize: number | undefined; classifierContextPerTurnChars: number | undefined; classifierContextIncludeAssistantTurns: boolean | undefined; + sessionAffinity: boolean; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; semanticMatchingEnabled: boolean; @@ -35,6 +36,7 @@ export interface ComplexityRouterConfigPayload { classifier_context_window_size?: number; classifier_context_per_turn_chars?: number; classifier_context_include_assistant_turns?: boolean; + session_affinity: boolean; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -78,6 +80,7 @@ export const buildComplexityRouterConfig = ({ classifierContextWindowSize, classifierContextPerTurnChars, classifierContextIncludeAssistantTurns, + sessionAffinity, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, @@ -110,6 +113,7 @@ export const buildComplexityRouterConfig = ({ classifierContextIncludeAssistantTurns !== undefined && { classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns, }), + session_affinity: sessionAffinity, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index f8d46f9ddd68..971c833a0de8 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -199,3 +199,28 @@ describe("buildUpdatedComplexityRouterConfig assistant turns", () => { expect(result.classifier_context_include_assistant_turns).toBeUndefined(); }); }); + +describe("buildUpdatedComplexityRouterConfig session affinity", () => { + it("writes session_affinity=false when the toggle is off", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity: false }); + expect(result.session_affinity).toBe(false); + }); + + it("writes session_affinity=true when the toggle is on", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity: true }); + expect(result.session_affinity).toBe(true); + }); + + it("re-asserts the backend's off-by-default when the form value is absent, rather than dropping the key", () => { + const result = buildUpdatedComplexityRouterConfig({ ...STORED, session_affinity: true }, FORM_VALUE); + expect(result.session_affinity).toBe(false); + }); + + it("stops a stored session_affinity=true from surviving a save that turned the toggle back off", () => { + const result = buildUpdatedComplexityRouterConfig( + { ...STORED, session_affinity: true }, + { ...FORM_VALUE, session_affinity: false }, + ); + expect(result.session_affinity).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index 17fa810b5295..eb5bb46f0e88 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -47,6 +47,7 @@ const expectedClassifiedTierConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + session_affinity: false, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, adaptive_eligible: "classified_tier", @@ -66,6 +67,7 @@ const expectedAdaptiveDisabledConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + session_affinity: false, }; describe("buildUpdatedComplexityRouterConfig", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 98b7ac519f55..c0806befa526 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -243,3 +243,76 @@ describe("EditAutoRouterModal assistant turns", () => { expect(savedConfig().classifier_context_include_assistant_turns).toBe(false); }); }); + +describe("EditAutoRouterModal session affinity", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const renderWithStoredConfig = (complexity_router_config: Record) => + renderWithProviders( + , + ); + + // A stored config with no session_affinity key now runs with affinity OFF, because the backend + // field defaults to False. The toggle has to render what the router actually does, and an + // untouched save must not flip it. + it("shows a stored config with no session_affinity key as off", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Session Affinity")); + expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(false); + }); + + it("shows a stored session_affinity=true as on and preserves it through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); + + await user.click(await screen.findByText("Advanced: Session Affinity")); + expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(true); + }); + + it("persists turning session affinity on", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(true); + }); + + it("persists turning session affinity back off", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); + + await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 99b5ff178b12..a70fc31d6fe3 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -13,6 +13,7 @@ import { hydrateKeywordTierRules, serializeKeywordTierRules } from "../add_model import ComplexityRouterConfig, { ComplexityRouterConfigValue, DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_SESSION_AFFINITY, DEFAULT_TIER_DISTANCE_PENALTY, } from "../add_model/ComplexityRouterConfig"; import NotificationsManager from "../molecules/notifications_manager"; @@ -36,6 +37,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_context_window_size", "classifier_context_per_turn_chars", "classifier_context_include_assistant_turns", + "session_affinity", "adaptive", "adaptive_weights", "tier_distance_penalty", @@ -101,6 +103,7 @@ export const buildUpdatedComplexityRouterConfig = ( value.classifier_context_include_assistant_turns !== undefined && { classifier_context_include_assistant_turns: value.classifier_context_include_assistant_turns, }), + session_affinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, ...(customTechnicalKeywords && customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords, @@ -218,6 +221,10 @@ const EditAutoRouterModal: React.FC = ({ typeof parsedConfig.classifier_context_include_assistant_turns === "boolean" ? parsedConfig.classifier_context_include_assistant_turns : undefined, + session_affinity: + typeof parsedConfig.session_affinity === "boolean" + ? parsedConfig.session_affinity + : DEFAULT_SESSION_AFFINITY, adaptive: parsedConfig.adaptive || false, adaptive_weights: parsedConfig.adaptive_weights, tier_distance_penalty: parsedConfig.tier_distance_penalty,