From af2b66bb517fac6a39b800977d7d0919278d669d Mon Sep 17 00:00:00 2001
From: Abhinav Kumar <abhinavkumar130503@gmail.com>
Date: Sat, 26 Jul 2025 11:32:34 +0530
Subject: [PATCH] feat: Add support for Bedrock custom inference profiles via
 model_id

---
 docs/docs/usage-guide/changing_a_model.md     | 20 +++++++++++++++++++
 .../algo/ai_handlers/litellm_ai_handler.py    |  5 +++++
 pr_agent/settings/.secrets_template.toml      |  1 +
 pr_agent/settings/configuration.toml          |  1 +
 4 files changed, 27 insertions(+)

diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md
index 361abbca..9c232d4e 100644
--- a/docs/docs/usage-guide/changing_a_model.md
+++ b/docs/docs/usage-guide/changing_a_model.md
@@ -250,6 +250,26 @@ model="bedrock/us.meta.llama4-scout-17b-instruct-v1:0"
 fallback_models=["bedrock/us.meta.llama4-maverick-17b-instruct-v1:0"]
 ```
 
+#### Custom Inference Profiles
+
+To use a custom inference profile with Amazon Bedrock (for cost allocation tags and other configuration settings), add the `model_id` parameter to your configuration:
+
+```toml
+[config] # in configuration.toml
+model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
+fallback_models=["bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"]
+
+[aws]
+AWS_ACCESS_KEY_ID="..."
+AWS_SECRET_ACCESS_KEY="..."
+AWS_REGION_NAME="..."
+
+[litellm]
+model_id = "your-custom-inference-profile-id"
+```
+
+The `model_id` parameter will be passed to all Bedrock completion calls, allowing you to use custom inference profiles for better cost allocation and reporting.
+
 See [litellm](https://docs.litellm.ai/docs/providers/bedrock#usage) documentation for more information about the environment variables required for Amazon Bedrock.
 
 ### DeepSeek
diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
index 63e9aaa1..a22906fe 100644
--- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py
+++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -352,6 +352,11 @@ class LiteLLMAIHandler(BaseAiHandler):
             # Support for custom OpenAI body fields (e.g., Flex Processing)
             kwargs = _process_litellm_extra_body(kwargs)
 
+            # Support for Bedrock custom inference profile via model_id
+            if get_settings().get("LITELLM.MODEL_ID", None) and 'bedrock/' in model:
+                kwargs["model_id"] = get_settings().litellm.model_id
+                get_logger().info(f"Using Bedrock custom inference profile: {get_settings().litellm.model_id}")
+
             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 c3d7a3f9..4e7d156a 100644
--- a/pr_agent/settings/.secrets_template.toml
+++ b/pr_agent/settings/.secrets_template.toml
@@ -19,6 +19,7 @@ key = ""  # Acquire through https://platform.openai.com
 # OpenAI Flex Processing (optional, for cost savings)
 # [litellm]
 # extra_body='{"processing_mode": "flex"}'
+# model_id = "" # Optional: Custom inference profile ID for Amazon Bedrock
 
 [pinecone]
 api_key = "..."
diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml
index 37ffbb7a..80bfc639 100644
--- a/pr_agent/settings/configuration.toml
+++ b/pr_agent/settings/configuration.toml
@@ -334,6 +334,7 @@ enable_callbacks = false
 success_callback = []
 failure_callback = []
 service_callback = []
+# model_id = "" # Optional: Custom inference profile ID for Amazon Bedrock
 
 [pr_similar_issue]
 skip_comments = false