Skip to content

Commit

Permalink
chore: improve overal stability of ChatLiteLLMComponent
Browse files Browse the repository at this point in the history
  • Loading branch information
berrytern committed Jun 27, 2024
1 parent db24a64 commit e3e5dae
Showing 1 changed file with 9 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -128,10 +128,14 @@ def build_model(self) -> LanguageModel:
raise ChatLiteLLMException(
"Could not import litellm python package. " "Please install it with `pip install litellm`"
)

# Set the API key based on the provider
self.kwargs[self.provider] = self.api_key
self.model_kwargs["api_key"] = self.api_key
# Remove empty keys
if "" in self.kwargs:
del self.kwargs[""]
if "" in self.model_kwargs:
del self.model_kwargs[""]
# Set the API key
# self.model_kwargs["api_key"] = self.api_key
# Report missing fields for Azure provider
if self.provider == "Azure":
if "api_base" not in self.kwargs:
raise Exception("Missing api_base on kwargs")
Expand All @@ -150,5 +154,6 @@ def build_model(self) -> LanguageModel:
max_retries=self.max_retries,
**self.kwargs,
)
llm.client.api_key = self.api_key

return llm

0 comments on commit e3e5dae

Please sign in to comment.