xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-03 12:20:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1564 from chandan84/fix/support_litellm_extra_headers

Browse Source 
Fix/support litellm extra headers
This commit is contained in:
Tal
2025-02-26 10:15:22 +02:00
committed by GitHub
parent 25ba9414fe c89c0eab8c
commit d6f405dd0d
2 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
pr_agent/algo/ai_handlers/litellm_ai_handler.py
View File

@ -11,6 +11,7 @@ from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
from pr_agent.algo.utils import ReasoningEffort, get_version
from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger
import json
OPENAI_RETRIES = 5
@ -254,12 +255,22 @@ class LiteLLMAIHandler(BaseAiHandler):
if self.repetition_penalty:
kwargs["repetition_penalty"] = self.repetition_penalty
#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):
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
get_logger().debug("Prompts", artifact={"system": system, "user": user})
if get_settings().config.verbosity_level >= 2:
get_logger().info(f"\nSystem prompt:\n{system}")
get_logger().info(f"\nUser prompt:\n{user}")
response = await acompletion(**kwargs)
except (openai.APIError, openai.APITimeoutError) as e:
get_logger().warning(f"Error during LLM inference: {e}")