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

Update litellm_ai_handler.py

Browse Source 
line 253-258, pass extra_headers fields from settings to litellm, exception handling to check if extra_headers is in dict format
This commit is contained in:
chandan84
2025-02-22 14:43:03 -05:00
committed by GitHub
parent 0e4a1d9ab8
commit 71451de156
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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: