Proposal: Complexity Based Model Routing

[bentito]

I created a PR, praxis-proxy/ai#446, just to see if this idea might work well within Praxis and to share the idea with @rikatz
I've mentioned this concept to @shaneutt in passing as well, seeing where this might best be solved. It might fit under the umbrella of topics in praxis-proxy/ai#74 or it might not be solvable with Praxis today, I'm not sure.

But backing up a bit, my interest is to create something like the model router found in the gemini-cli coding agent. That tool relied on a light LLM to judge complexity of the prompt for a given turn of a conversation with the agent and route it to the best matching LLM in its list.

Obviously solving this on the wire has nice benefits in that it could be used with any agent and, if done right, should just give better solutions cheaper and faster with no complications on the agent side. If a given prompt is judged complex it heads off to a bigger, slower, more complex model for a better response, if it's easy it heads to light, fast, cheap model.

Sounds straightforward, but one issue I've realized as I think about it more is that in a multi-turn conversation how do you maintain state? When this concept is implemented within an agent it's more obvious, the agent is maintaining state via its memory functions. If this is implemented in Praxis, it's an important problem to solve, answers from various LLM without the context being tied together will not work. I'm not sure how you solve it, but I think it can be solved such that the model routing mechanism can keep an agent's complete conversation rational despite multiple LLMs being in the loop.

I'm sure there are some links I could add above, please advise, b/c I know there are already efforts in Praxis on cost based routing as well as maintaining state.

[cnuland]

My personal research on this has shown significant value in operating off the turn-level embeddings while still routing the entire context. Most failures I've seen with the implementation of other open-source solutions, like vLLM Semantic Router, occur when semantic routing is applied to the entire context. Embedding models compress a text span into a fixed-dimensional vector by effectively averaging/pooling semantic signal across all tokens. As context length grows, the embedding becomes a diluted composite of many topics rather than a sharp representation of current intent, leading to semantic dilution, and it degrades discriminative power for similarity-based routing

Adding to this, there are significant challenges of shared context. If model A is local and handles PPI, and model B is external and can't handle PPI due to compliance, there's a high risk that the router will share sensitive data across an egress security boundary.

I'm more than willing to assist with this! Let me know how I can help.

[franciscojavierarceo]

Most failures I've seen with the implementation of other open-source solutions, like vLLM Semantic Router, occur when semantic routing is applied to the entire context. Embedding models compress a text span into a fixed-dimensional vector by effectively averaging/pooling semantic signal across all tokens. As context length grows, the embedding becomes a diluted composite of many topics rather than a sharp representation of current intent, leading to semantic dilution, and it degrades discriminative power for similarity-based routing

+1 this would be a generic problem with any aggregation over an input so isolated queries is probably one approach to handle that.

[bentito]

Wow that PPI usage and wrinkle is a doozy! 😄 Yeah my note below that there this is a storage layer on the horizon should help but will need to be tempered by that consideration.

I agree per turn is the way to make the model choice decision. A vector on something like complexity scores drawn from the whole conversation would not tell me what model to choose per turn. Think that very complex work got done early on, then we're finishing a whole solution with git commit, git push -- unless some else pulled down that vector it's still likely to tell us we need a heavy model for the simple tasks.

[bentito]

Thinking about this further, maybe help for multi-turn continuity exists in the form of some proposals in the works now

1. Request-Scoped State (Proposal #432: ResponsesState/Extensions (https://github.com/praxis-proxy/ai/blob/main/docs/proposals/00432_responses-state.md)): This introduces a generic, type-safe extension container on the HttpFilterContext

2. Cross-Request Persistence (Proposal #412: Storage Layer (https://github.com/praxis-proxy/ai/blob/main/docs/proposals/00412_storage_layer.md)): To solve the multi-turn continuity problem, the arbiter needs to remember previous reasoning steps and tool outputs across multiple HTTP cycles. This proposal defines pluggable, distributed backend storage (KV and Object) for exactly this purpose.

[leseb]

I think 2 is actually https://github.com/praxis-proxy/praxis/blob/main/docs/proposals/00786_iterative-request-router.md?

[caldeirav]

I think it is clear that the proxy architecture must adopt some kind of stateful, multi-layered session control plane to avoid typical issues such as misclassification due to context collapse but also beyond reasoning failures, the potential model bouncing which introduces financial penalties and tool execution loops.

An example of implementation in vLLM-SR is Conversational Routing Momentum (check vllm-project/semantic-router#1458) which use asymmetric attack/release dynamics to track complexity trends over time, paired with a cache-aware switch gate that locks session affinity unless generation savings on a smaller model explicitly exceed the uncached re-hydration cost. I'd suggest we can start with an inventory of approaches / techniques including research literature if it can help but ultimately going with some real-world use cases and being able to test / evolve such decision algorithm at control plane level will be critical to test / validate which approach are relevant.

Two things I'd like to highlight in the design approach:

• Whatever processing / routing design we use in the control plane, we need to have a quick fail-safe mechanism typically being able to catch when a model's output fails parameter validation or triggers a tool exception, so that we have an immediate re-escalation to a higher model tier.
• Finally, having transparent classification scores and routing execution with human-interpretable skill-mapping vectors and ability to expose explicit control primitives is a key dependency to be able to build an evaluation capability for this type of deliverable (but also provides the auditability needed by our customers).

[usize]

tl;dr I think we can make real progress on mixture of models by treating classification as a black box and just building the machinery to safely do rewrites.

Existing epic is praxis-proxy/ai#74 -- I can post a comment there. Thoughts follow.

The classification portion of this is one discussion and nuances around things like context pollution are important for specific implementations, but, for Praxis specifically, I think the most important thing we can do is reason about what the overall flow of model selection looks like.

Here's what I mean. We can treat model selection as a black box and solve the general problems in the space. This will let operators experiment with whatever classification schemes fit their needs best. As an infra provider, I think that's our appropriate scope, right?

The generic problem in the space comes down to enforcing security boundaries and compliance policies while also allowing model rewriting.

I think the philosophical stance of Praxis is separation of concerns via filter composition whenever possible.

Here I think it looks like:

classify -> validate and select -> rewrite

Where each one can be its own filter.

The classifier can promote some headers with the common signals we'd expect: complexity, category, etc...

The selection filter can be cluster aware, and select models that are available to the request based on its properties and our configuration. It can also provide validation against per model policies like, maybe an openrouter model isn't allowed to have PII. Our classifier could output something like a Praxis-Context-Has-PII: true.

Then finally, we can promote our eventual model to a header like Praxis-Model-Rewrite and a rewrite filter can actually do the work of swapping out the model field. Potentially in a way that doesn't require buffering the whole request. I think this fits as a separate filter because we might have per-provider logic that helps us e.g., rewrite a chunk of a streaming request to avoid buffering-- risky but doable IMO.

@franciscojavierarceo @bentito @cnuland

I was thinking of proposing something like this "generic model selection" filter set, where the classifier is a black box with a contract based on well known headers or our state object, and then a general purpose select/validate and model rewrite filter can work across any.

How would that sit with you all's model of the world?

[franciscojavierarceo]

I was thinking of proposing something like this "generic model selection" filter set, where the classifier is a black box with a contract based on well known headers or our state object, and then a general purpose select/validate and model rewrite filter can work across any.

I like this approach. Probably we need to have some configurable mapping between the model_id and the classification and the big brain LLM can choose the the routing dynamically based on some prompt.

I don't think rewrite is a requirement though, because I imagine in some cases rewrite/query expansion could result in unexpected consequences. So probably we can make that configurable.

[usize]

When you say rewrite/query expansion, what do you mean?

I mean replacing the model field in the body. e.g., opus-4.8 rewritten to kimi-k3 or something along those lines.

Nothing more extensive than that, aside from the mechanism and the safety checks we'd want in place to ensure that the payload still has integrity.

I only brought it up because I want to treat model picking as a black box, since there are lots of different interesting approaches to mapping context to models and I'd like Praxis to have some common gate in front of slotting them in.

[franciscojavierarceo]

When you say rewrite/query expansion, what do you mean?

Ah, my mistake, I misinterpreted what this. Probably from my brain over-indexing on the Responses API rewrite_query is where my head went).

I mean replacing the model field in the body. e.g., opus-4.8 rewritten to kimi-k3 or something along those lines.

Makes sense, then I'm definitely aligned with this. 👍

[bentito]

I was thinking of proposing something like this "generic model selection" filter set, where the classifier is a black box with a contract based on well known headers or our state object, and then a general purpose select/validate and model rewrite filter can work across any.

Thanks for helping framing this in a better way for actual implementation in Praxis!

Is "our state object" already a thing or how does that relate to "Cross-Request Persistence (Proposal #412: Storage Layer"? I want to clarify because I think it all falls apart fast if Praxis can't keep a multi-turn conversation coherent across several LLMs. It really in many ways has to do the work of a sub-agent in order to do that.

I'm also not sure about rewrite. I was thinking more like classic routing with the agent's message intact. I think trying to rewrite for PII might end badly and for things like complexity score, it really is just a question of what level LLM you send the message to.

[franciscojavierarceo]

FWIW, PII Sanitization I think is separate from rewrite IMO as the similarity between the two for sanitization should be high (even according to something like Levenshtein's distance) compared to query rewrite which may not need be the case at all (which could dwarf intent, in positive or negative directions).

[usize]

When I brought up rewriting, I only meant modifying the model field in the request body. Not mutating the entire payload.

Cross-Request persistence is out of scope for what I had in mind, which is just abstracting out the basic machinery for consuming the output of some model picker (anything that maps a request to a model) to the safety checks and model field rewriting that should be common across all approaches to model picking.

By state object I meant ctx.filter_metadata or something similar, in lieu of promoting a header.

[leseb]

@bentito check ctx.extensions and IRR filter

[bentito]

@leseb this is awesome! b/c this is where I think this whole discussion is heading. praxis needs to act as an agent to do the model routing properly. I'll try to rework my earlier praxis-proxy/ai#446 (that would never really work due to lack of state knowledge of the client agent conversation) to use ctx.extensions and the IRR filter