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

line 253-255, pass extra_headers fields from settings to litellm

Browse Source
This commit is contained in:
chandan84
2025-02-22 14:12:39 -05:00
parent 01cd66f4f5
commit e7b05732f8
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
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 get_version from pr_agent.algo.utils import get_version
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger from pr_agent.log import get_logger
import json
OPENAI_RETRIES = 5 OPENAI_RETRIES = 5
@ -249,6 +250,10 @@ class LiteLLMAIHandler(BaseAiHandler):
get_logger().info(f"\nSystem prompt:\n{system}") get_logger().info(f"\nSystem prompt:\n{system}")
get_logger().info(f"\nUser prompt:\n{user}") 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
if get_settings().get("LITELLM.EXTRA_HEADERS", None):
kwargs["extra_headers"] = json.loads(get_settings().litellm.extra_headers)
response = await acompletion(**kwargs) response = await acompletion(**kwargs)
except (openai.APIError, openai.APITimeoutError) as e: except (openai.APIError, openai.APITimeoutError) as e:
get_logger().warning(f"Error during LLM inference: {e}") get_logger().warning(f"Error during LLM inference: {e}")