KeyArgo
/
smolagents
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add support for additional keyword arguments in LiteLLMModel

This commit is contained in:
chakib belgaid 2025-01-01 11:54:35 -08:00
parent 5991206ae5
commit 5fd88e5db7
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/smolagents/models.py
View File

@ -415,6 +415,7 @@ class LiteLLMModel(Model):
model_id="anthropic/claude-3-5-sonnet-20240620", model_id="anthropic/claude-3-5-sonnet-20240620",
api_base=None, api_base=None,
api_key=None, api_key=None,
**kwargs
): ):
super().__init__() super().__init__()
self.model_id = model_id self.model_id = model_id
@ -422,6 +423,7 @@ class LiteLLMModel(Model):
litellm.add_function_to_prompt = True litellm.add_function_to_prompt = True
self.api_base = api_base self.api_base = api_base
self.api_key = api_key self.api_key = api_key
self.kwargs = kwargs
def __call__( def __call__(
self, self,
@ -441,6 +443,7 @@ class LiteLLMModel(Model):
max_tokens=max_tokens, max_tokens=max_tokens,
api_base=self.api_base, api_base=self.api_base,
api_key=self.api_key, api_key=self.api_key,
**self.kwargs,
) )
self.last_input_token_count = response.usage.prompt_tokens self.last_input_token_count = response.usage.prompt_tokens
self.last_output_token_count = response.usage.completion_tokens self.last_output_token_count = response.usage.completion_tokens
@ -465,6 +468,7 @@ class LiteLLMModel(Model):
max_tokens=max_tokens, max_tokens=max_tokens,
api_base=self.api_base, api_base=self.api_base,
api_key=self.api_key, api_key=self.api_key,
**self.kwargs,
) )
tool_calls = response.choices[0].message.tool_calls[0] tool_calls = response.choices[0].message.tool_calls[0]
self.last_input_token_count = response.usage.prompt_tokens self.last_input_token_count = response.usage.prompt_tokens