diff --git a/cookbook/fusion_routing/README.md b/cookbook/fusion_routing/README.md new file mode 100644 index 00000000000..d08e83064e3 --- /dev/null +++ b/cookbook/fusion_routing/README.md @@ -0,0 +1,46 @@ +# Fusion-style dual-model routing + +Cognition's [Devin Fusion](https://cognition.com/blog/devin-fusion) pairs a frontier "main" model with a cheap "sidekick" model inside one agentic session: the main model plans, interprets ambiguity, and reviews; the sidekick handles mechanical, well-specified subtasks. A classifier decides per-call which one should run, and both stay live so a mid-session switch doesn't cost a fresh cache warm-up. + +This cookbook shows the pattern on LiteLLM Proxy using a plain `CustomLogger.async_pre_call_hook`. No core LiteLLM changes needed. + +## How it differs from `auto_router` + +LiteLLM already ships an [`auto_router`](https://docs.litellm.ai/docs/proxy/auto_routing) with a semantic router (embedding similarity against example utterances) and a [complexity router](https://docs.litellm.ai/docs/proxy/auto_routing#complexity-router) (keyword/token heuristics). Both pick one model per request from the text of the current prompt. + +Fusion routing is a different signal: it's session-aware (every call in a task shares a `session_id`) and it classifies by *role in the workflow* (main vs. sidekick), not by the text's apparent complexity. A short "apply this diff" call can still be a main-model judgment call, and a long code-heavy call can still be pure sidekick work. Nothing stops you from wiring `complexity_router`'s scoring in as one signal inside your `FusionClassifier` implementation below; this cookbook just adds the session/role layer on top of whatever classifier you plug in. + +## Files + +- `fusion_hook.py` — `FusionRoutingHook`, a `CustomLogger` that rewrites `data["model"]` based on a pluggable `FusionClassifier`. Ships with `ToolNameFusionClassifier`, a placeholder heuristic that routes by the last tool call name. +- `config.example.yaml` — proxy config wiring the hook in via `litellm_settings.callbacks`. +- `test_fusion_hook.py` — unit tests for the routing/metadata behavior. + +## Quick start + +1. Define your main + sidekick deployments and register the hook in `config.yaml` (see `config.example.yaml`). +2. Write a `FusionClassifier` for your workflow. `ToolNameFusionClassifier` is a starting point — swap in whatever signal fits (tool name, a cheap LLM-as-judge call, `complexity_router` scoring, etc). +3. Start the proxy: + + ```bash + litellm --config cookbook/fusion_routing/config.example.yaml + ``` + +4. Send every call in one task through the same `litellm_session_id` so the hook can see them as one session: + + ```bash + curl http://localhost:4000/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "fusion-main", + "litellm_session_id": "task-42", + "messages": [{"role": "user", "content": "read config.yaml and summarize it"}] + }' + ``` + + The response's `model` field (and `metadata.fusion_role` in the logged request) shows which deployment actually served the call. Filter `/ui/?page=logs` or your spend export by `session_id` to see the blended main/sidekick cost for a task versus an all-frontier baseline. + +## Caveat + +`ToolNameFusionClassifier` is intentionally dumb — it exists to show the hook's shape, not to be a production classifier. The real payoff comes from a classifier that understands your agent's workflow (what's a plan/interpret/review step vs. a mechanical execution step), not from text-surface heuristics. diff --git a/cookbook/fusion_routing/config.example.yaml b/cookbook/fusion_routing/config.example.yaml new file mode 100644 index 00000000000..f5c99d1bab8 --- /dev/null +++ b/cookbook/fusion_routing/config.example.yaml @@ -0,0 +1,13 @@ +model_list: + - model_name: fusion-main + litellm_params: + model: anthropic/claude-opus-4-8 + api_key: os.environ/ANTHROPIC_API_KEY + - model_name: fusion-sidekick + litellm_params: + model: anthropic/claude-haiku-4-5-20251001 + api_key: os.environ/ANTHROPIC_API_KEY + +litellm_settings: + callbacks: + - fusion_hook.fusion_hook_instance diff --git a/cookbook/fusion_routing/fusion_hook.py b/cookbook/fusion_routing/fusion_hook.py new file mode 100644 index 00000000000..caf443087dc --- /dev/null +++ b/cookbook/fusion_routing/fusion_hook.py @@ -0,0 +1,73 @@ +from dataclasses import dataclass +from typing import Any, Dict, FrozenSet, List, Literal, Optional, Protocol, Union + +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import get_session_id_from_request_data +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + +FusionRole = Literal["main", "sidekick"] + + +@dataclass(frozen=True, slots=True) +class FusionModelGroup: + main_model: str + sidekick_model: str + + def model_for_role(self, role: FusionRole) -> str: + return self.main_model if role == "main" else self.sidekick_model + + +class FusionClassifier(Protocol): + async def classify(self, data: Dict[str, Any], session_id: str) -> FusionRole: ... + + +def _last_tool_call_name(data: Dict[str, Any]) -> Optional[str]: + messages: List[Dict[str, Any]] = data.get("messages") or [] + for message in reversed(messages): + tool_calls: List[Dict[str, Any]] = message.get("tool_calls") or [] + for tool_call in tool_calls: + return tool_call.get("function", {}).get("name") + return None + + +class ToolNameFusionClassifier: + def __init__(self, sidekick_tools: FrozenSet[str]) -> None: + self._sidekick_tools = sidekick_tools + + async def classify(self, data: Dict[str, Any], session_id: str) -> FusionRole: + tool_name = _last_tool_call_name(data) + return "sidekick" if tool_name in self._sidekick_tools else "main" + + +class FusionRoutingHook(CustomLogger): + def __init__(self, group: FusionModelGroup, classifier: FusionClassifier) -> None: + super().__init__() + self._group = group + self._classifier = classifier + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: Dict[str, Any], + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, Dict[str, Any]]]: + session_id = get_session_id_from_request_data(data) + if session_id is None: + return None + + role = await self._classifier.classify(data, session_id) + metadata: Dict[str, Any] = { + **(data.get("metadata") or {}), + "fusion_role": role, + "fusion_original_model": data["model"], + } + return {**data, "model": self._group.model_for_role(role), "metadata": metadata} + + +fusion_hook_instance = FusionRoutingHook( + group=FusionModelGroup(main_model="fusion-main", sidekick_model="fusion-sidekick"), + classifier=ToolNameFusionClassifier(sidekick_tools=frozenset({"read_file", "grep", "run_tests", "apply_diff"})), +) diff --git a/cookbook/fusion_routing/test_fusion_hook.py b/cookbook/fusion_routing/test_fusion_hook.py new file mode 100644 index 00000000000..2f0c0990233 --- /dev/null +++ b/cookbook/fusion_routing/test_fusion_hook.py @@ -0,0 +1,100 @@ +from typing import Any, Dict, FrozenSet, Optional, cast + +import pytest + +from cookbook.fusion_routing.fusion_hook import ( + FusionModelGroup, + FusionRole, + FusionRoutingHook, + ToolNameFusionClassifier, +) +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth + +USER_API_KEY_DICT = UserAPIKeyAuth() +CACHE = DualCache() + + +class StubClassifier: + def __init__(self, role: FusionRole) -> None: + self._role = role + + async def classify(self, data: Dict[str, Any], session_id: str) -> FusionRole: + return self._role + + +def _make_hook(role: FusionRole) -> FusionRoutingHook: + group = FusionModelGroup(main_model="fusion-main", sidekick_model="fusion-sidekick") + return FusionRoutingHook(group=group, classifier=StubClassifier(role)) + + +async def _route(hook: FusionRoutingHook, data: Dict[str, Any]) -> Optional[Dict[str, Any]]: + result = await hook.async_pre_call_hook(USER_API_KEY_DICT, CACHE, data, "completion") + return cast(Optional[Dict[str, Any]], result) + + +async def test_routes_to_sidekick_model(): + hook = _make_hook("sidekick") + data = {"model": "fusion-main", "litellm_session_id": "task-1", "messages": []} + + result = await _route(hook, data) + + assert result is not None + assert result["model"] == "fusion-sidekick" + assert result["metadata"]["fusion_role"] == "sidekick" + assert result["metadata"]["fusion_original_model"] == "fusion-main" + + +async def test_routes_to_main_model(): + hook = _make_hook("main") + data = {"model": "fusion-sidekick", "litellm_session_id": "task-1", "messages": []} + + result = await _route(hook, data) + + assert result is not None + assert result["model"] == "fusion-main" + assert result["metadata"]["fusion_role"] == "main" + + +async def test_noop_without_session_id(): + hook = _make_hook("sidekick") + data = {"model": "fusion-main", "messages": []} + + result = await _route(hook, data) + + assert result is None + + +async def test_preserves_existing_metadata(): + hook = _make_hook("sidekick") + data = { + "model": "fusion-main", + "litellm_session_id": "task-1", + "messages": [], + "metadata": {"user_id": "u-1"}, + } + + result = await _route(hook, data) + + assert result is not None + assert result["metadata"]["user_id"] == "u-1" + assert result["metadata"]["fusion_role"] == "sidekick" + + +@pytest.mark.parametrize( + "sidekick_tools,tool_name,expected_role", + [ + (frozenset({"read_file", "grep"}), "read_file", "sidekick"), + (frozenset({"read_file", "grep"}), "edit_file", "main"), + (frozenset({"read_file", "grep"}), None, "main"), + ], +) +async def test_tool_name_classifier( + sidekick_tools: FrozenSet[str], tool_name: Optional[str], expected_role: FusionRole +): + classifier = ToolNameFusionClassifier(sidekick_tools=sidekick_tools) + messages = [{"role": "assistant", "tool_calls": [{"function": {"name": tool_name}}]}] if tool_name else [] + + role = await classifier.classify({"messages": messages}, session_id="task-1") + + assert role == expected_role