Support litellm router features?

[benwhalley]

litellm includes some nice features to rate limit calls, and switch between deployments based on https://docs.litellm.ai/docs/routing

I wonder if it would be possible to configure magentic to accept a router config?

e.g. on line 12 here, rather than just importing litellm and then using litellm.completion you could create a router instance based on a MAGENTIC_LITELLM_ROUTER config and below on line 61 or 94 use router.completion or router.acompletion rather than litellm.acompletion?

https://github.com/jackmpcollins/magentic/blob/5b4b1012e630bd01baa0ef8bf3caf86fb68ae993/src/magentic/chat_model/litellm_chat_model.py

My initial use case is just setting RPM limits, but I can see it would also be nice to be able to track token usage.

[jackmpcollins]

It's not currently possible but it could definitely be added. The code you link is old. The newer version would be even more simple to update.

magentic/src/magentic/chat_model/litellm_chat_model.py

Lines 153 to 155 in 6747d99

	response = discard_none_arguments(litellm.completion)(
	model=self.model,

I think a good approach would be to add a router: litellm.Router | None = None parameter to LitellmChatModel.__init__. Then in (a)complete use router.completion instead of litellm.completion if router was provided (assuming they meet the same interface, which appears to be the case).

[benwhalley]

would you be open to a pr if I have a go at this?

[jackmpcollins]

would you be open to a pr if I have a go at this?

@benwhalley Absolutely!