From e46ef9d6429f339fbff4cc197caa3a90d1e3156e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 11 Jun 2025 22:19:17 -0700 Subject: [PATCH] test: update test with new kwargs --- tests/local_testing/test_custom_llm.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_custom_llm.py b/tests/local_testing/test_custom_llm.py index 77f4544afa..61312f4e52 100644 --- a/tests/local_testing/test_custom_llm.py +++ b/tests/local_testing/test_custom_llm.py @@ -265,8 +265,10 @@ class MyCustomLLM(CustomLLM): print_verbose: Callable, logging_obj: Any, optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, litellm_params=None, - aembedding=None, ) -> EmbeddingResponse: model_response.model = model @@ -289,6 +291,9 @@ class MyCustomLLM(CustomLLM): print_verbose: Callable, logging_obj: Any, optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, litellm_params=None, ) -> EmbeddingResponse: model_response.model = model