Add support for additional keyword arguments in LiteLLMModel
This commit is contained in:
		
							parent
							
								
									5991206ae5
								
							
						
					
					
						commit
						5fd88e5db7
					
				|  | @ -415,6 +415,7 @@ class LiteLLMModel(Model): | |||
|         model_id="anthropic/claude-3-5-sonnet-20240620", | ||||
|         api_base=None, | ||||
|         api_key=None, | ||||
|         **kwargs | ||||
|     ): | ||||
|         super().__init__() | ||||
|         self.model_id = model_id | ||||
|  | @ -422,6 +423,7 @@ class LiteLLMModel(Model): | |||
|         litellm.add_function_to_prompt = True | ||||
|         self.api_base = api_base | ||||
|         self.api_key = api_key | ||||
|         self.kwargs = kwargs | ||||
| 
 | ||||
|     def __call__( | ||||
|         self, | ||||
|  | @ -441,6 +443,7 @@ class LiteLLMModel(Model): | |||
|             max_tokens=max_tokens, | ||||
|             api_base=self.api_base, | ||||
|             api_key=self.api_key, | ||||
|             **self.kwargs, | ||||
|         ) | ||||
|         self.last_input_token_count = response.usage.prompt_tokens | ||||
|         self.last_output_token_count = response.usage.completion_tokens | ||||
|  | @ -465,6 +468,7 @@ class LiteLLMModel(Model): | |||
|             max_tokens=max_tokens, | ||||
|             api_base=self.api_base, | ||||
|             api_key=self.api_key, | ||||
|             **self.kwargs, | ||||
|         ) | ||||
|         tool_calls = response.choices[0].message.tool_calls[0] | ||||
|         self.last_input_token_count = response.usage.prompt_tokens | ||||
|  |  | |||
		Loading…
	
		Reference in New Issue