search

jackmpcollins / magentic

Seamlessly integrate LLMs as Python functions
https://magentic.dev/
MIT License
1.95k stars 96 forks source link

Support litellm router features? #314

Open benwhalley opened 1 month ago

benwhalley commented 1 month ago

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 commented 1 month ago

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.

https://github.com/jackmpcollins/magentic/blob/6747d99dc37f0becbfbb12fdfdc17392da343ee0/src/magentic/chat_model/litellm_chat_model.py#L153-L155

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 commented 1 month ago

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

jackmpcollins commented 1 month ago

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

@benwhalley Absolutely!