Add LiteLLM Router model #1096

NishantBaheti · 2025-03-28T15:26:16Z

Simply add a LiteLLM Router to the Model Pool

Motivation - When we talk about agentic frameworks LLM's are always a bottleneck in terms of API calls per minute. What if we simply add LiteLLM router in smolagents.

Guidelines - https://github.com/huggingface/smolagents/blob/main/CONTRIBUTING.md

Code Snippet -

import os
from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMRouter

os.environ["OPENAI_API_KEY"] = ""
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION"] = ""

llm_loadbalancer_model_list = [
    {
        "model_name": "model-group-1",
        "litellm_params": {
            "model": "gpt-4o-mini",
            "api_key": os.getenv("OPENAI_API_KEY"),
        },
    },
    {
        "model_name": "model-group-1",
        "litellm_params": {
            "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
            "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
            "aws_region_name": os.getenv("AWS_REGION"),
        },
    }
]


model = LiteLLMRouter(
    model_id="model-group-1",
    model_list=llm_loadbalancer_model_list,
    router_kwargs={"routing_strategy": "simple-shuffle"},
)
agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model)

agent.run("How many seconds would it take for a leopard at full speed to run through Pont des Arts?")

@aymeric-roucher Let me know your thoughts. Excited to be a part of this community.

luiztauffer · 2025-04-03T11:55:51Z

+1 to that! definitely needed!

NishantBaheti · 2025-04-03T16:33:56Z

@luiztauffer Idea is to give user the possibility to provide model's list from config like etcd etc and the group name of the models to choose from the list of models. Hence user can manage multiple models with different combinations in different use cases like text2sql, websearch, summarizes etc. I have put one example as well. model_id is actually the group of model that you want to use from the list. May be I should change the name to model group, just wanted to follow the same naming conventions. 😊

luiztauffer · 2025-04-03T16:43:24Z

@NishantBaheti yeah I agree. I missed that point at the beginning, but I got it right after trying your code =)
Which works, btw! I'm already using it in my projects, thanks for the contribution! I hope they review it soon.

NishantBaheti · 2025-04-03T17:00:58Z

Hopefully, Felt like this should be here, it can help a lot of engineers trying to scale agents for a million users. BTW loved this framework, it was easy to use and easy to contribute.

luiztauffer · 2025-04-04T11:11:02Z

two suggestions:

It looks like the method create_client must be defined by classes inheriting from ApiModel:

def create_client(self):
        """Create a client for the model."""
        try:
            from litellm import Router
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "Please install 'litellm' extra to use LiteLLMModel: `pip install 'smolagents[litellm]'`"
            )
        self.router = Router(
            model_list=self._model_list,
            **self.router_config,
        )

It would be useful to pass forward extra configurations for the litellm router:

router_config = {
    "routing_strategy": "simple-shuffle",
    "num_retries": 3,
    "retry_after": 10,
}
model = LiteLLMRouter(model_id="all_models", model_list=model_list, router_config=router_config)

class LiteLLMRouter(ApiModel):
    def __init__(
        self,
        model_id: str,
        model_list: List[Dict[Any, Any]],
        router_config: Optional[Dict[str, Any]] = None,
        ...

NishantBaheti · 2025-04-04T15:23:42Z

@luiztauffer Good Suggestions. Added. Check now.

NishantBaheti · 2025-04-05T10:21:07Z

@aymeric-roucher WDYT?

aymeric-roucher · 2025-04-11T13:42:43Z

src/smolagents/models.py

+            )
+        self.router: Router = Router(model_list=self._model_list, **self._router_kwargs)
+
+    def __call__(


Since this __call__method is the same as LiteLLMModel, could you just have this LiteLLMRouter class inherit from LiteLLMModel , and only define init and create_client methods?

aymeric-roucher · 2025-04-11T13:43:18Z

@NishantBaheti this will be a great addition, thanks a lot! Just left some comments.

aymeric-roucher · 2025-04-11T13:44:53Z

src/smolagents/models.py

+            if flatten_messages_as_text is not None
+            else self.model_id.startswith(("ollama", "groq", "cerebras"))
+        )
+        self.create_client()


You should modify the base create_client methods to accept kwargs, and pass it router_kwargs or {} : thus you can get rid of self._router_kwargs and self._model_list attributes.

aymeric-roucher · 2025-04-11T13:45:18Z

src/smolagents/models.py

@@ -905,6 +905,132 @@ def __call__(
        return self.postprocess_message(first_message, tools_to_call_from)


+class LiteLLMRouter(ApiModel):


Have this directly inherit from LiteLLMModel to get its call method.

NishantBaheti · 2025-04-17T11:46:43Z

@aymeric-roucher Made some changes based on the review. Let me know

NishantBaheti · 2025-04-20T12:10:09Z

albertvillanova

Thanks for the update. I think there are still requested changes that haven't been addressed yet.

NishantBaheti · 2025-04-22T16:29:47Z

@albertvillanova Thanks a lot for the response!!!! Yeah, I missed one. But it was a repeated one. As I've already mentioned in the comment LiteLLM model completion and router completion are two different methods and, according to my understanding, cannot be used interchangeably. Is there a better way available?

albertvillanova

Thanks for the follow-up!

You are right that the completions may differ conceptually, but in this case, I believe basic OOP method overriding should be sufficient. You can have LiteLLMRouter inherit from LiteLLMModel, and simply override the create_client and __init__ methods as needed, while still leveraging the shared logic in __call__.

This would help reduce duplication and make the design more maintainable. Let me know if you run into any issues with that approach or if there is a specific part that does not fit this pattern.

Also, feel free to add your comments in the reviewer comments you have above (this link: #1096 (comment)), so we can follow the discussion: I cannot find your comment

As I've already mentioned in the comment

…S concept

NishantBaheti · 2025-04-22T18:24:13Z

@albertvillanova
Solved this now,

Old strucuture

ApiModel
- LiteLLMModel
- LiteLLMRouter

New Structure

ApiModel
- LiteLLMModel
  - LiteLLMRouter

Please have a look and let me know.. I believe all the asks are completed :)

HuggingFaceDocBuilderDev · 2025-04-23T08:08:12Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

albertvillanova

Thanks for your contribution.

I made a few small adjustments to fix typos and align the code style with the rest of the codebase.

NishantBaheti · 2025-04-23T09:12:37Z

@albertvillanova Yes, I got it. I understand we all love our codebase as it is our creation, like a painting. Even I want some things in certain ways, following a certain structure, and IT NEEDS TO BE PERFECT—the forbidden curse of perfection.
Anyway, I'm always happy to help. Again, kudos on the library—it's easy to read, easy to implement, and easy to contribute to. I have more feature suggestions, and might contribute in the future too. :)

adding LiteLLM Router As Well :)

fee55d0

adding test, updating doc and modifying code a little bit

10258a9

style correction

da1fc02

adding create client separately and adding router configurations options

71e6988

style change

8ad1962

aymeric-roucher reviewed Apr 11, 2025

View reviewed changes

review changes

4a30f55

NishantBaheti requested a review from aymeric-roucher April 12, 2025 10:05

Merge branch 'main' into adding-litellmrouter

829ab98

albertvillanova requested changes Apr 22, 2025

View reviewed changes

albertvillanova reviewed Apr 22, 2025

View reviewed changes

NishantBaheti added 3 commits April 22, 2025 17:41

inheriting LiteLLMModel class for LiteLLMRouter class to maintain OOP…

e765838

…S concept

inheriting LiteLLMModel class for LiteLLMRouter class to maintain OOP…

0b520c9

…S concept

styling

f7984b3

NishantBaheti requested a review from albertvillanova April 22, 2025 18:28

Fix typo in tests

9689457

albertvillanova added 16 commits April 22, 2025 21:57

Fix error message

8c6c5e3

Make model_id required, since default is deprecated

5043096

Fix typo in docstring

76d2635

Fix typo in docstring

265e905

Fix type hints

b2ca003

Fix docstring

38d1635

Simplify test

b9da886

Rename class

27fbe7b

Refactor test

af7b9af

Test create_client

68dd360

Fix docstring

1003a59

Improve docstring

bb64098

Fix class name in error message

8b7113e

Simplify init and create_client

510760e

Rename router_kwargs to client_kwargs

acf1755

Add LiteLLMRouterModel to Reference docs

a3a19e3

albertvillanova approved these changes Apr 23, 2025

View reviewed changes

albertvillanova changed the title ~~adding LiteLLM Router (LLM Pooling) As Well~~ Add LiteLLM Router model Apr 23, 2025

albertvillanova merged commit 6923eb5 into huggingface:main Apr 23, 2025
4 checks passed

NishantBaheti deleted the adding-litellmrouter branch April 23, 2025 09:56

		@@ -905,6 +905,132 @@ def __call__(
		return self.postprocess_message(first_message, tools_to_call_from)


		class LiteLLMRouter(ApiModel):

Add LiteLLM Router model #1096

Add LiteLLM Router model #1096

Uh oh!

Conversation

NishantBaheti commented Mar 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Simply add a LiteLLM Router to the Model Pool

Uh oh!

luiztauffer commented Apr 3, 2025

Uh oh!

NishantBaheti commented Apr 3, 2025

Uh oh!

luiztauffer commented Apr 3, 2025

Uh oh!

NishantBaheti commented Apr 3, 2025

Uh oh!

luiztauffer commented Apr 4, 2025

Uh oh!

NishantBaheti commented Apr 4, 2025

Uh oh!

NishantBaheti commented Apr 5, 2025

Uh oh!

aymeric-roucher Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

aymeric-roucher commented Apr 11, 2025

Uh oh!

aymeric-roucher Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

aymeric-roucher Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

NishantBaheti commented Apr 17, 2025

Uh oh!

NishantBaheti commented Apr 20, 2025

Uh oh!

albertvillanova left a comment

Choose a reason for hiding this comment

Uh oh!

NishantBaheti commented Apr 22, 2025

Uh oh!

albertvillanova left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NishantBaheti commented Apr 22, 2025

Uh oh!

HuggingFaceDocBuilderDev commented Apr 23, 2025

Uh oh!

albertvillanova left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

NishantBaheti commented Apr 23, 2025

Uh oh!

Uh oh!

NishantBaheti commented Mar 28, 2025 •

edited

Loading

albertvillanova left a comment •

edited

Loading