Fix/support litellm extra headers

[chandan84]

User description

Description

With reference to issue - #1557

This pull request modifies the litellm_ai_handler.py script to allow passing extra_headers when making requests through LiteLLM. This is particularly useful when using an API management gateway that requires specific headers for security and authorization.

Changes:

Modified: pr_agent/algo/ai_handlers/litellm_ai_handler.py

Lines 253-261:

if get_settings().get("LITELLM.EXTRA_HEADERS", None):
        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

Purpose:

This change introduces the ability to pass extra_headers in pr-agent config for LiteLLM AI Client via the LITELLM.EXTRA_HEADERS setting. The value of this setting should be a JSON string representing the desired headers.

This enables users to pass custom headers, such as authorization tokens or API keys, when routing requests through an API management gateway.

---

PR Type

Enhancement, Bug fix

---

Description

• Added support for extra_headers in LiteLLM requests.

• Validated LITELLM.EXTRA_HEADERS as a JSON object.

• Introduced error handling for invalid extra_headers configuration.

• Enabled passing custom headers for security and authorization.

---

Changes walkthrough 📝

	Relevant files
Enhancement	litellm_ai_handler.py Support and validate extra headers for LiteLLM                      pr_agent/algo/ai_handlers/litellm_ai_handler.py Added logic to handle LITELLM.EXTRA_HEADERS from settings. Validated extra_headers as a JSON object. Introduced exception handling for invalid JSON or non-dict formats. Enabled passing custom headers via kwargs["extra_headers"]. +11/-0

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-merge-pro-for-open-source]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1557 - Partially compliant Compliant requirements: Support passing custom headers to Azure OpenAI via litellm Non-compliant requirements: Allow setting headers through environment variables Requires further human verification: Verify that the fix resolves the specific 'projectId' missing header error with Azure OpenAI

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Error Handling The error handling for invalid JSON could be improved by providing more specific error messages and guidance to users on the expected format 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)}") Configuration Validation The code should validate that required header fields are present in the provided JSON object before passing it to litellm 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")

[qodo-merge-pro-for-open-source]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Cache parsed headers for performance Move the JSON parsing of LITELLM.EXTRA_HEADERS outside the chat_completion function to avoid parsing the same JSON on every request. Cache the parsed headers. pr_agent/algo/ai_handlers/litellm_ai_handler.py [254-261] -if get_settings().get("LITELLM.EXTRA_HEADERS", None): - 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 +# At class level +_cached_litellm_headers = None +# In __init__ or a setup method +def _setup_litellm_headers(self): + if get_settings().get("LITELLM.EXTRA_HEADERS", None): + try: + headers = json.loads(get_settings().litellm.extra_headers) + if not isinstance(headers, dict): + raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") + self._cached_litellm_headers = headers + except json.JSONDecodeError as e: + raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") + +# In chat_completion +if self._cached_litellm_headers: + kwargs["extra_headers"] = self._cached_litellm_headers + Apply this suggestion Suggestion importance[1-10]: 7 __ Why: The suggestion improves performance by avoiding redundant JSON parsing on every request, and properly caches the headers at the class level. This is a meaningful optimization for code that may be called frequently.	Medium
Learned best practice	Enhance error handling and input validation for configuration parsing to prevent potential runtime errors The code should handle the case where litellm.extra_headers might be None even after the initial settings check. Also, consider adding more specific error handling for the JSON parsing to provide clearer error messages. pr_agent/algo/ai_handlers/litellm_ai_handler.py [254-261] -if get_settings().get("LITELLM.EXTRA_HEADERS", None): +extra_headers = get_settings().get("LITELLM.EXTRA_HEADERS") +if extra_headers: try: - litellm_extra_headers = json.loads(get_settings().litellm.extra_headers) + litellm_extra_headers = json.loads(extra_headers) if not isinstance(litellm_extra_headers, dict): - raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object") + raise ValueError("LITELLM.EXTRA_HEADERS must be a JSON object, got {type(litellm_extra_headers)}") + kwargs["extra_headers"] = litellm_extra_headers except json.JSONDecodeError as e: - raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") - kwargs["extra_headers"] = litellm_extra_headers + get_logger().error(f"Failed to parse LITELLM.EXTRA_HEADERS: {str(e)}") + raise ValueError(f"LITELLM.EXTRA_HEADERS must be a valid JSON object: {str(e)}") Apply this suggestion Suggestion importance[1-10]: 6	Low
More

• Author self-review: I have reviewed the PR code suggestions, and addressed the relevant ones.