Merge pull request #470 from kikumoto/feature/update_aws_bedrock_claude_implementation

Feature/update aws bedrock claude implementation

Author: tjbck

examples/pipelines/providers/aws_bedrock_claude_pipeline.py

@@ -12,7 +12,7 @@
 import json
 import logging
 from io import BytesIO
-from typing import List, Union, Generator, Iterator
+from typing import List, Union, Generator, Iterator, Optional, Any
 
 import boto3
 
@@ -23,12 +23,23 @@
 
 from utils.pipelines.main import pop_system_message
 
+REASONING_EFFORT_BUDGET_TOKEN_MAP = {
+    "none": None,
+    "low": 1024,
+    "medium": 4096,
+    "high": 16384,
+    "max": 32768,
+}
+
+# Maximum combined token limit for Claude 3.7
+MAX_COMBINED_TOKENS = 64000
+
 
 class Pipeline:
     class Valves(BaseModel):
-        AWS_ACCESS_KEY: str = ""
-        AWS_SECRET_KEY: str = ""
-        AWS_REGION_NAME: str = ""
+        AWS_ACCESS_KEY: Optional[str] = None
+        AWS_SECRET_KEY: Optional[str] = None
+        AWS_REGION_NAME: Optional[str] = None
 
     def __init__(self):
         self.type = "manifold"
@@ -47,21 +58,25 @@ def __init__(self):
             }
         )
 
-        self.bedrock = boto3.client(aws_access_key_id=self.valves.AWS_ACCESS_KEY,
-                                    aws_secret_access_key=self.valves.AWS_SECRET_KEY,
-                                    service_name="bedrock",
-                                    region_name=self.valves.AWS_REGION_NAME)
-        self.bedrock_runtime = boto3.client(aws_access_key_id=self.valves.AWS_ACCESS_KEY,
-                                            aws_secret_access_key=self.valves.AWS_SECRET_KEY,
-                                            service_name="bedrock-runtime",
-                                            region_name=self.valves.AWS_REGION_NAME)
+        self.valves = self.Valves(
+            **{
+                "AWS_ACCESS_KEY": os.getenv("AWS_ACCESS_KEY", ""),
+                "AWS_SECRET_KEY": os.getenv("AWS_SECRET_KEY", ""),
+                "AWS_REGION_NAME": os.getenv(
+                    "AWS_REGION_NAME", os.getenv(
+                        "AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "")
+                    )
+                ),
+            }
+        )
 
-        self.pipelines = self.get_models()
+        self.update_pipelines()
 
 
     async def on_startup(self):
         # This function is called when the server is started.
         print(f"on_startup:{__name__}")
+        self.update_pipelines()
         pass
 
     async def on_shutdown(self):
@@ -72,40 +87,58 @@ async def on_shutdown(self):
     async def on_valves_updated(self):
         # This function is called when the valves are updated.
         print(f"on_valves_updated:{__name__}")
-        self.bedrock = boto3.client(aws_access_key_id=self.valves.AWS_ACCESS_KEY,
-                                    aws_secret_access_key=self.valves.AWS_SECRET_KEY,
-                                    service_name="bedrock",
-                                    region_name=self.valves.AWS_REGION_NAME)
-        self.bedrock_runtime = boto3.client(aws_access_key_id=self.valves.AWS_ACCESS_KEY,
-                                            aws_secret_access_key=self.valves.AWS_SECRET_KEY,
-                                            service_name="bedrock-runtime",
-                                            region_name=self.valves.AWS_REGION_NAME)
-        self.pipelines = self.get_models()
-
-    def pipelines(self) -> List[dict]:
-        return self.get_models()
+        self.update_pipelines()
+
+    def update_pipelines(self) -> None:
+        try:
+            self.bedrock = boto3.client(service_name="bedrock",
+                                        aws_access_key_id=self.valves.AWS_ACCESS_KEY,
+                                        aws_secret_access_key=self.valves.AWS_SECRET_KEY,
+                                        region_name=self.valves.AWS_REGION_NAME)
+            self.bedrock_runtime = boto3.client(service_name="bedrock-runtime",
+                                                aws_access_key_id=self.valves.AWS_ACCESS_KEY,
+                                                aws_secret_access_key=self.valves.AWS_SECRET_KEY,
+                                                region_name=self.valves.AWS_REGION_NAME)
+            self.pipelines = self.get_models()
+        except Exception as e:
+            print(f"Error: {e}")
+            self.pipelines = [
+                {
+                    "id": "error",
+                    "name": "Could not fetch models from Bedrock, please set up AWS Key/Secret or Instance/Task Role.",
+                },
+            ]
 
     def get_models(self):
-        if self.valves.AWS_ACCESS_KEY and self.valves.AWS_SECRET_KEY:
-            try:
-                response = self.bedrock.list_foundation_models(byProvider='Anthropic', byInferenceType='ON_DEMAND')
-                return [
-                    {
-                        "id": model["modelId"],
-                        "name": model["modelName"],
-                    }
-                    for model in response["modelSummaries"]
-                ]
-            except Exception as e:
-                print(f"Error: {e}")
-                return [
-                    {
-                        "id": "error",
-                        "name": "Could not fetch models from Bedrock, please update the Access/Secret Key in the valves.",
-                    },
-                ]
-        else:
-            return []
+        try:
+            res = []
+            response = self.bedrock.list_foundation_models(byProvider='Anthropic')
+            for model in response['modelSummaries']:
+                inference_types = model.get('inferenceTypesSupported', [])
+                if "ON_DEMAND" in inference_types:
+                    res.append({'id': model['modelId'], 'name': model['modelName']})
+                elif "INFERENCE_PROFILE" in inference_types:
+                    inferenceProfileId = self.getInferenceProfileId(model['modelArn'])
+                    if inferenceProfileId:
+                        res.append({'id': inferenceProfileId, 'name': model['modelName']})
+
+            return res
+        except Exception as e:
+            print(f"Error: {e}")
+            return [
+                {
+                    "id": "error",
+                    "name": "Could not fetch models from Bedrock, please check permissoin.",
+                },
+            ]
+
+    def getInferenceProfileId(self, modelArn: str) -> str:
+        response = self.bedrock.list_inference_profiles()
+        for profile in response.get('inferenceProfileSummaries', []):
+            for model in profile.get('models', []):
+                if model.get('modelArn') == modelArn:
+                    return profile['inferenceProfileId']
+        return None
 
     def pipe(
         self, user_message: str, model_id: str, messages: List[dict], body: dict
@@ -139,11 +172,53 @@ def pipe(
 
             payload = {"modelId": model_id,
                        "messages": processed_messages,
-                       "system": [{'text': system_message if system_message else 'you are an intelligent ai assistant'}],
-                       "inferenceConfig": {"temperature": body.get("temperature", 0.5)},
-                       "additionalModelRequestFields": {"top_k": body.get("top_k", 200), "top_p": body.get("top_p", 0.9)}
+                       "system": [{'text': system_message["content"] if system_message else 'you are an intelligent ai assistant'}],
+                       "inferenceConfig": {
+                           "temperature": body.get("temperature", 0.5),
+                           "topP": body.get("top_p", 0.9),
+                           "maxTokens": body.get("max_tokens", 4096),
+                           "stopSequences": body.get("stop", []),
+                        },
+                        "additionalModelRequestFields": {"top_k": body.get("top_k", 200)}
                        }
+
             if body.get("stream", False):
+                supports_thinking = "claude-3-7" in model_id
+                reasoning_effort = body.get("reasoning_effort", "none")
+                budget_tokens = REASONING_EFFORT_BUDGET_TOKEN_MAP.get(reasoning_effort)
+
+                # 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 = None
+
+                if supports_thinking and budget_tokens:
+                    # Check if the combined tokens (budget_tokens + max_tokens) exceeds the limit
+                    max_tokens = payload.get("max_tokens", 4096)
+                    combined_tokens = budget_tokens + max_tokens
+
+                    if combined_tokens > MAX_COMBINED_TOKENS:
+                        error_message = f"Error: Combined tokens (budget_tokens {budget_tokens} + max_tokens {max_tokens} = {combined_tokens}) exceeds the maximum limit of {MAX_COMBINED_TOKENS}"
+                        print(error_message)
+                        return error_message
+
+                    payload["inferenceConfig"]["maxTokens"] = combined_tokens
+                    payload["additionalModelRequestFields"]["thinking"] = {
+                        "type": "enabled",
+                        "budget_tokens": budget_tokens,
+                    }
+                    # Thinking requires temperature 1.0 and does not support top_p, top_k
+                    payload["inferenceConfig"]["temperature"] = 1.0
+                    if "top_k" in payload["additionalModelRequestFields"]:
+                        del payload["additionalModelRequestFields"]["top_k"]
+                    if "topP" in payload["inferenceConfig"]:
+                        del payload["inferenceConfig"]["topP"]
                 return self.stream_response(model_id, payload)
             else:
                 return self.get_completion(model_id, payload)
@@ -152,30 +227,45 @@ def pipe(
 
     def process_image(self, image: str):
         img_stream = None
+        content_type = None
+
         if image["url"].startswith("data:image"):
-            if ',' in image["url"]:
-                base64_string = image["url"].split(',')[1]
+            mime_type, base64_string = image["url"].split(",", 1)
+            content_type = mime_type.split(":")[1].split(";")[0]
             image_data = base64.b64decode(base64_string)
-
             img_stream = BytesIO(image_data)
         else:
-            img_stream = requests.get(image["url"]).content
+            response = requests.get(image["url"])
+            img_stream = BytesIO(response.content)
+            content_type = response.headers.get('Content-Type', 'image/jpeg')
+
+        media_type = content_type.split('/')[-1] if '/' in content_type else content_type
         return {
-            "image": {"format": "png" if image["url"].endswith(".png") else "jpeg",
-                      "source": {"bytes": img_stream.read()}}
+            "image": {
+                "format": media_type,
+                "source": {"bytes": img_stream.read()}
+            }
         }
 
     def stream_response(self, model_id: str, payload: dict) -> Generator:
-        if "system" in payload:
-            del payload["system"]
-        if "additionalModelRequestFields" in payload:
-            del payload["additionalModelRequestFields"]
         streaming_response = self.bedrock_runtime.converse_stream(**payload)
+
+        in_resasoning_context = False
         for chunk in streaming_response["stream"]:
-            if "contentBlockDelta" in chunk:
-                yield chunk["contentBlockDelta"]["delta"]["text"]
+            if in_resasoning_context and "contentBlockStop" in chunk:
+                in_resasoning_context = False
+                yield "\n </think> \n\n"
+            elif "contentBlockDelta" in chunk and "delta" in chunk["contentBlockDelta"]:
+                if "reasoningContent" in chunk["contentBlockDelta"]["delta"]:
+                    if not in_resasoning_context:
+                        yield "<think>"
+
+                    in_resasoning_context = True
+                    if "text" in chunk["contentBlockDelta"]["delta"]["reasoningContent"]:
+                        yield chunk["contentBlockDelta"]["delta"]["reasoningContent"]["text"]
+                elif "text" in chunk["contentBlockDelta"]["delta"]:
+                    yield chunk["contentBlockDelta"]["delta"]["text"]
 
     def get_completion(self, model_id: str, payload: dict) -> str:
         response = self.bedrock_runtime.converse(**payload)
         return response['output']['message']['content'][0]['text']
-