From 84983f3e9d753f04f2af42228bb416a05cea75b5 Mon Sep 17 00:00:00 2001 From: chandan84 Date: Sat, 22 Feb 2025 14:56:17 -0500 Subject: [PATCH] line 253-261, pass extra_headers fields from settings to litellm, exception handling to check if extra_headers is in dict format --- pr_agent/algo/ai_handlers/litellm_ai_handler.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 016a5eb3..56125aff 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -252,9 +252,12 @@ class LiteLLMAIHandler(BaseAiHandler): #Added support for extra_headers while using litellm to call underlying model, via a api management gateway, would allow for passing custom headers for security and authorization if get_settings().get("LITELLM.EXTRA_HEADERS", None): - litellm_extra_headers = json.loads(get_settings().litellm.extra_headers) - if not isinstance(litellm_extra_headers, dict): - raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") + try: + litellm_extra_headers = json.loads(get_settings().litellm.extra_headers) + if not isinstance(litellm_extra_headers, dict): + raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") + except json.JSONDecodeError as e: + raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") kwargs["extra_headers"] = litellm_extra_headers response = await acompletion(**kwargs)