From 12af211c13450a025e106afd1a10296c17b899fa Mon Sep 17 00:00:00 2001 From: Abhinav Kumar Date: Mon, 7 Jul 2025 21:14:45 +0530 Subject: [PATCH 1/5] feat: support OpenAI Flex Processing via [litellm] extra_body config --- docs/docs/usage-guide/changing_a_model.md | 10 ++++++++++ pr_agent/algo/ai_handlers/litellm_ai_handler.py | 10 ++++++++++ pr_agent/settings/.secrets_template.toml | 4 ++++ 3 files changed, 24 insertions(+) diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md index 9648e6cf..361abbca 100644 --- a/docs/docs/usage-guide/changing_a_model.md +++ b/docs/docs/usage-guide/changing_a_model.md @@ -32,6 +32,16 @@ OPENAI__API_BASE=https://api.openai.com/v1 OPENAI__KEY=sk-... ``` +### OpenAI Flex Processing + +To reduce costs for non-urgent/background tasks, enable Flex Processing: + +```toml +[litellm] +extra_body='{"processing_mode": "flex"}' +``` + +See [OpenAI Flex Processing docs](https://platform.openai.com/docs/guides/flex-processing) for details. ### Azure diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index ec96d952..ecb84ea7 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -364,6 +364,16 @@ class LiteLLMAIHandler(BaseAiHandler): raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") kwargs["extra_headers"] = litellm_extra_headers + # Support for custom OpenAI body fields (e.g., Flex Processing) + if get_settings().get("LITELLM.EXTRA_BODY", None): + try: + litellm_extra_body = json.loads(get_settings().litellm.extra_body) + if not isinstance(litellm_extra_body, dict): + raise ValueError("LITELLM.EXTRA_BODY must be a JSON object") + kwargs.update(litellm_extra_body) + except json.JSONDecodeError as e: + raise ValueError(f"LITELLM.EXTRA_BODY contains invalid JSON: {str(e)}") + get_logger().debug("Prompts", artifact={"system": system, "user": user}) if get_settings().config.verbosity_level >= 2: diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml index 350abe5c..c3d7a3f9 100644 --- a/pr_agent/settings/.secrets_template.toml +++ b/pr_agent/settings/.secrets_template.toml @@ -16,6 +16,10 @@ key = "" # Acquire through https://platform.openai.com #deployment_id = "" # The deployment name you chose when you deployed the engine #fallback_deployments = [] # For each fallback model specified in configuration.toml in the [config] section, specify the appropriate deployment_id +# OpenAI Flex Processing (optional, for cost savings) +# [litellm] +# extra_body='{"processing_mode": "flex"}' + [pinecone] api_key = "..." environment = "gcp-starter" From 6a55bbcd2332278e277888a0d2e5111ecc2e25d5 Mon Sep 17 00:00:00 2001 From: Abhinav Kumar Date: Mon, 7 Jul 2025 21:20:25 +0530 Subject: [PATCH 2/5] fix: prevent LITELLM.EXTRA_BODY from overriding existing parameters in LiteLLMAIHandler --- pr_agent/algo/ai_handlers/litellm_ai_handler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index ecb84ea7..14d98546 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -365,11 +365,14 @@ class LiteLLMAIHandler(BaseAiHandler): kwargs["extra_headers"] = litellm_extra_headers # Support for custom OpenAI body fields (e.g., Flex Processing) - if get_settings().get("LITELLM.EXTRA_BODY", None): + if get_settings().litellm.extra_body: try: litellm_extra_body = json.loads(get_settings().litellm.extra_body) if not isinstance(litellm_extra_body, dict): raise ValueError("LITELLM.EXTRA_BODY must be a JSON object") + colliding_keys = kwargs.keys() & litellm_extra_body.keys() + if colliding_keys: + raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}") kwargs.update(litellm_extra_body) except json.JSONDecodeError as e: raise ValueError(f"LITELLM.EXTRA_BODY contains invalid JSON: {str(e)}") From 8127d52ab3035bfc0c8358a133e94f66ee8c38a6 Mon Sep 17 00:00:00 2001 From: Abhinav Kumar Date: Mon, 7 Jul 2025 21:26:13 +0530 Subject: [PATCH 3/5] fix: security checks --- pr_agent/algo/ai_handlers/litellm_ai_handler.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 14d98546..cececc60 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -365,11 +365,17 @@ class LiteLLMAIHandler(BaseAiHandler): kwargs["extra_headers"] = litellm_extra_headers # 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: - litellm_extra_body = json.loads(get_settings().litellm.extra_body) + litellm_extra_body = json.loads(extra_body) if not isinstance(litellm_extra_body, dict): 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() if colliding_keys: raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}") From e2d71acb9d28ed26282c3196c9f5b34cb880c59b Mon Sep 17 00:00:00 2001 From: Abhinav Kumar Date: Mon, 7 Jul 2025 21:27:35 +0530 Subject: [PATCH 4/5] fix: remove comments --- pr_agent/algo/ai_handlers/litellm_ai_handler.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index cececc60..2dd4ad31 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -365,7 +365,6 @@ class LiteLLMAIHandler(BaseAiHandler): kwargs["extra_headers"] = litellm_extra_headers # Support for custom OpenAI body fields (e.g., Flex Processing) - # 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: From e0d70837682df90dbf61cd32df620cff27742d98 Mon Sep 17 00:00:00 2001 From: Abhinav Kumar Date: Wed, 9 Jul 2025 12:04:26 +0530 Subject: [PATCH 5/5] feat: refactor LITELLM.EXTRA_BODY processing into a dedicated method --- .../algo/ai_handlers/litellm_ai_handler.py | 48 ++++++++++++------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 2dd4ad31..cbfe37da 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -175,6 +175,37 @@ class LiteLLMAIHandler(BaseAiHandler): response_log['main_pr_language'] = 'unknown' return response_log + def _process_litellm_extra_body(self, kwargs: dict) -> dict: + """ + Process LITELLM.EXTRA_BODY configuration and update kwargs accordingly. + + Args: + kwargs: The current kwargs dictionary to update + + Returns: + Updated kwargs dictionary + + Raises: + ValueError: If extra_body contains invalid JSON, unsupported keys, or colliding keys + """ + allowed_extra_body_keys = {"processing_mode", "service_tier"} + extra_body = getattr(getattr(get_settings(), "litellm", None), "extra_body", None) + if extra_body: + try: + litellm_extra_body = json.loads(extra_body) + if not isinstance(litellm_extra_body, dict): + 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() + if colliding_keys: + raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}") + kwargs.update(litellm_extra_body) + except json.JSONDecodeError as e: + raise ValueError(f"LITELLM.EXTRA_BODY contains invalid JSON: {str(e)}") + return kwargs + def _configure_claude_extended_thinking(self, model: str, kwargs: dict) -> dict: """ Configure Claude extended thinking parameters if applicable. @@ -365,22 +396,7 @@ class LiteLLMAIHandler(BaseAiHandler): kwargs["extra_headers"] = litellm_extra_headers # Support for custom OpenAI body fields (e.g., Flex Processing) - allowed_extra_body_keys = {"processing_mode", "service_tier"} - extra_body = getattr(getattr(get_settings(), "litellm", None), "extra_body", None) - if extra_body: - try: - litellm_extra_body = json.loads(extra_body) - if not isinstance(litellm_extra_body, dict): - 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() - if colliding_keys: - raise ValueError(f"LITELLM.EXTRA_BODY cannot override existing parameters: {', '.join(colliding_keys)}") - kwargs.update(litellm_extra_body) - except json.JSONDecodeError as e: - raise ValueError(f"LITELLM.EXTRA_BODY contains invalid JSON: {str(e)}") + kwargs = self._process_litellm_extra_body(kwargs) get_logger().debug("Prompts", artifact={"system": system, "user": user})