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

fix: security checks

Browse Source
This commit is contained in:
Abhinav Kumar
2025-07-07 21:26:13 +05:30
parent 6a55bbcd23
commit 8127d52ab3
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -365,11 +365,17 @@ class LiteLLMAIHandler(BaseAiHandler):
kwargs["extra_headers"] = litellm_extra_headers kwargs["extra_headers"] = litellm_extra_headers
# Support for custom OpenAI body fields (e.g., Flex Processing) # Support for custom OpenAI body fields (e.g., Flex Processing)
if get_settings().litellm.extra_body: # Only allow whitelisted keys for security
allowed_extra_body_keys = {"processing_mode", "service_tier"}
extra_body = getattr(getattr(get_settings(), "litellm", None), "extra_body", None)
if extra_body:
try: try:
litellm_extra_body = json.loads(get_settings().litellm.extra_body) litellm_extra_body = json.loads(extra_body)
if not isinstance(litellm_extra_body, dict): if not isinstance(litellm_extra_body, dict):
raise ValueError("LITELLM.EXTRA_BODY must be a JSON object") raise ValueError("LITELLM.EXTRA_BODY must be a JSON object")
unsupported_keys = set(litellm_extra_body.keys()) - allowed_extra_body_keys
if unsupported_keys:
raise ValueError(f"LITELLM.EXTRA_BODY contains unsupported keys: {', '.join(unsupported_keys)}. Allowed keys: {', '.join(allowed_extra_body_keys)}")
colliding_keys = kwargs.keys() & litellm_extra_body.keys() colliding_keys = kwargs.keys() & litellm_extra_body.keys()
if colliding_keys: if colliding_keys:
raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}") raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}")