From c11a144a9036b180c6f9395ba95520e7b6246fa6 Mon Sep 17 00:00:00 2001 From: Aviraj <100823015+avirajsingh7@users.noreply.github.com> Date: Mon, 1 Sep 2025 19:35:27 +0530 Subject: [PATCH] fix lite llm api key --- py_zerox/pyzerox/models/modellitellm.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/py_zerox/pyzerox/models/modellitellm.py b/py_zerox/pyzerox/models/modellitellm.py index e7b2d204..edbf5d6c 100644 --- a/py_zerox/pyzerox/models/modellitellm.py +++ b/py_zerox/pyzerox/models/modellitellm.py @@ -52,7 +52,12 @@ def system_prompt(self, prompt: str) -> None: ## custom method on top of BaseModel def validate_environment(self) -> None: """Validates the environment variables required for the model.""" - env_config = litellm.validate_environment(model=self.model) + env_config = litellm.validate_environment( + model=self.model, + api_key=self.kwargs.get("api_key"), + api_base=self.kwargs.get("api_base"), + api_version=self.kwargs.get("api_version"), + ) if not env_config["keys_in_environment"]: raise MissingEnvironmentVariables(extra_info=env_config) @@ -64,7 +69,7 @@ def validate_model(self) -> None: def validate_access(self) -> None: """Validates access to the model -> if environment variables are set correctly with correct values.""" - if not litellm.check_valid_key(model=self.model,api_key=None): + if not litellm.check_valid_key(model=self.model, api_key=self.kwargs.get("api_key")): raise ModelAccessError(extra_info={"model": self.model})