Adding budget values

claude_sonnet-3-7-bedrock.py

@@ -139,12 +139,32 @@ class Pipeline:
 
                 processed_messages.append({"role": message["role"], "content": processed_content})
 
-            reasoning_config = {
+            # Set budget tokens for reasoning
-                "thinking": {
+                reasoning_effort = body.get("reasoning_effort", "medium")
-                    "type": "enabled",
+                budget_tokens = REASONING_EFFORT_BUDGET_TOKEN_MAP.get(reasoning_effort)
-                    "budget_tokens": 4096
+
+                # Allow users to input an integer value representing budget tokens
+                if (
+                    not budget_tokens
+                    and reasoning_effort not in REASONING_EFFORT_BUDGET_TOKEN_MAP.keys()
+                ):
+                    try:
+                        budget_tokens = int(reasoning_effort)
+                    except ValueError as e:
+                        print("Failed to convert reasoning effort to int", e)
+                        budget_tokens = 4096
+
+            # Do not use thinking if budget_tokens is set to None
+
+            if budget_tokens is not None:
+                reasoning_config = {
+                    "thinking": {
+                        "type": "enabled",
+                        "budget_tokens": budget_tokens
+                    }
                 }
-            }
+            else:
+                reasoning_config = {}
             payload = {"modelId": model_id,
                        "messages": processed_messages,
                        "system": [{'text': system_message if system_message else 'you are an intelligent ai assistant'}],