From 71451de156f172114f9c237fcfd260adaa056f9b Mon Sep 17 00:00:00 2001
From: chandan84 <61125092+chandan84@users.noreply.github.com>
Date: Sat, 22 Feb 2025 14:43:03 -0500
Subject: [PATCH] Update litellm_ai_handler.py

line 253-258, 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, 5 insertions(+), 4 deletions(-)

diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
index 4ab5ee34..016a5eb3 100644
--- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py
+++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -251,10 +251,11 @@ class LiteLLMAIHandler(BaseAiHandler):
                 get_logger().info(f"\nUser prompt:\n{user}")
 
             #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
-            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")
-            kwargs["extra_headers"] = litellm_extra_headers
+            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")
+                kwargs["extra_headers"] = litellm_extra_headers
                 
             response = await acompletion(**kwargs)
         except (openai.APIError, openai.APITimeoutError) as e: