two version of R2R are hereHEADmaster

5 files changed, 665 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
new file mode 100644
index 00000000..de46edb9
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
@@ -0,0 +1,106 @@
+import types
+from typing import List, Optional
+
+from openai.types.image import Image
+
+from litellm.types.llms.bedrock import (
+    AmazonNovaCanvasTextToImageRequest, AmazonNovaCanvasTextToImageResponse,
+    AmazonNovaCanvasTextToImageParams, AmazonNovaCanvasRequestBase,
+)
+from litellm.types.utils import ImageResponse
+
+
+class AmazonNovaCanvasConfig:
+    """
+    Reference: https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/model-catalog/serverless/amazon.nova-canvas-v1:0
+
+    """
+
+    @classmethod
+    def get_config(cls):
+        return {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not k.startswith("__")
+               and not isinstance(
+                v,
+                (
+                    types.FunctionType,
+                    types.BuiltinFunctionType,
+                    classmethod,
+                    staticmethod,
+                ),
+            )
+               and v is not None
+        }
+
+    @classmethod
+    def get_supported_openai_params(cls, model: Optional[str] = None) -> List:
+        """
+        """
+        return ["n", "size", "quality"]
+
+    @classmethod
+    def _is_nova_model(cls, model: Optional[str] = None) -> bool:
+        """
+        Returns True if the model is a Nova Canvas model
+
+        Nova models follow this pattern:
+
+        """
+        if model:
+            if "amazon.nova-canvas" in model:
+                return True
+        return False
+
+    @classmethod
+    def transform_request_body(
+            cls, text: str, optional_params: dict
+    ) -> AmazonNovaCanvasRequestBase:
+        """
+        Transform the request body for Amazon Nova Canvas model
+        """
+        task_type = optional_params.pop("taskType", "TEXT_IMAGE")
+        image_generation_config = optional_params.pop("imageGenerationConfig", {})
+        image_generation_config = {**image_generation_config, **optional_params}
+        if task_type == "TEXT_IMAGE":
+            text_to_image_params = image_generation_config.pop("textToImageParams", {})
+            text_to_image_params = {"text" :text, **text_to_image_params}
+            text_to_image_params = AmazonNovaCanvasTextToImageParams(**text_to_image_params)
+            return AmazonNovaCanvasTextToImageRequest(textToImageParams=text_to_image_params, taskType=task_type,
+                                                      imageGenerationConfig=image_generation_config)
+        raise NotImplementedError(f"Task type {task_type} is not supported")
+
+    @classmethod
+    def map_openai_params(cls, non_default_params: dict, optional_params: dict) -> dict:
+        """
+        Map the OpenAI params to the Bedrock params
+        """
+        _size = non_default_params.get("size")
+        if _size is not None:
+            width, height = _size.split("x")
+            optional_params["width"], optional_params["height"] = int(width), int(height)
+        if non_default_params.get("n") is not None:
+            optional_params["numberOfImages"] = non_default_params.get("n")
+        if non_default_params.get("quality") is not None:
+            if non_default_params.get("quality") in ("hd", "premium"):
+                optional_params["quality"] = "premium"
+            if non_default_params.get("quality") == "standard":
+                optional_params["quality"] = "standard"
+        return optional_params
+
+    @classmethod
+    def transform_response_dict_to_openai_response(
+            cls, model_response: ImageResponse, response_dict: dict
+    ) -> ImageResponse:
+        """
+        Transform the response dict to the OpenAI response
+        """
+
+        nova_response = AmazonNovaCanvasTextToImageResponse(**response_dict)
+        openai_images: List[Image] = []
+        for _img in nova_response.get("images", []):
+            openai_images.append(Image(b64_json=_img))
+
+        model_response.data = openai_images
+        return model_response
diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability1_transformation.py b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability1_transformation.py
new file mode 100644
index 00000000..698ecca9
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability1_transformation.py
@@ -0,0 +1,104 @@
+import types
+from typing import List, Optional
+
+from openai.types.image import Image
+
+from litellm.types.utils import ImageResponse
+
+
+class AmazonStabilityConfig:
+    """
+    Reference: https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/providers?model=stability.stable-diffusion-xl-v0
+
+    Supported Params for the Amazon / Stable Diffusion models:
+
+    - `cfg_scale` (integer): Default `7`. Between [ 0 .. 35 ]. How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)
+
+    - `seed` (float): Default: `0`. Between [ 0 .. 4294967295 ]. Random noise seed (omit this option or use 0 for a random seed)
+
+    - `steps` (array of strings): Default `30`. Between [ 10 .. 50 ]. Number of diffusion steps to run.
+
+    - `width` (integer): Default: `512`. multiple of 64 >= 128. Width of the image to generate, in pixels, in an increment divible by 64.
+        Engine-specific dimension validation:
+
+        - SDXL Beta: must be between 128x128 and 512x896 (or 896x512); only one dimension can be greater than 512.
+        - SDXL v0.9: must be one of 1024x1024, 1152x896, 1216x832, 1344x768, 1536x640, 640x1536, 768x1344, 832x1216, or 896x1152
+        - SDXL v1.0: same as SDXL v0.9
+        - SD v1.6: must be between 320x320 and 1536x1536
+
+    - `height` (integer): Default: `512`. multiple of 64 >= 128. Height of the image to generate, in pixels, in an increment divible by 64.
+        Engine-specific dimension validation:
+
+        - SDXL Beta: must be between 128x128 and 512x896 (or 896x512); only one dimension can be greater than 512.
+        - SDXL v0.9: must be one of 1024x1024, 1152x896, 1216x832, 1344x768, 1536x640, 640x1536, 768x1344, 832x1216, or 896x1152
+        - SDXL v1.0: same as SDXL v0.9
+        - SD v1.6: must be between 320x320 and 1536x1536
+    """
+
+    cfg_scale: Optional[int] = None
+    seed: Optional[float] = None
+    steps: Optional[List[str]] = None
+    width: Optional[int] = None
+    height: Optional[int] = None
+
+    def __init__(
+        self,
+        cfg_scale: Optional[int] = None,
+        seed: Optional[float] = None,
+        steps: Optional[List[str]] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+    ) -> None:
+        locals_ = locals().copy()
+        for key, value in locals_.items():
+            if key != "self" and value is not None:
+                setattr(self.__class__, key, value)
+
+    @classmethod
+    def get_config(cls):
+        return {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not k.startswith("__")
+            and not isinstance(
+                v,
+                (
+                    types.FunctionType,
+                    types.BuiltinFunctionType,
+                    classmethod,
+                    staticmethod,
+                ),
+            )
+            and v is not None
+        }
+
+    @classmethod
+    def get_supported_openai_params(cls, model: Optional[str] = None) -> List:
+        return ["size"]
+
+    @classmethod
+    def map_openai_params(
+        cls,
+        non_default_params: dict,
+        optional_params: dict,
+    ):
+        _size = non_default_params.get("size")
+        if _size is not None:
+            width, height = _size.split("x")
+            optional_params["width"] = int(width)
+            optional_params["height"] = int(height)
+
+        return optional_params
+
+    @classmethod
+    def transform_response_dict_to_openai_response(
+        cls, model_response: ImageResponse, response_dict: dict
+    ) -> ImageResponse:
+        image_list: List[Image] = []
+        for artifact in response_dict["artifacts"]:
+            _image = Image(b64_json=artifact["base64"])
+            image_list.append(_image)
+
+        model_response.data = image_list
+
+        return model_response
diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability3_transformation.py b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability3_transformation.py
new file mode 100644
index 00000000..2c90b3a1
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/amazon_stability3_transformation.py
@@ -0,0 +1,100 @@
+import types
+from typing import List, Optional
+
+from openai.types.image import Image
+
+from litellm.types.llms.bedrock import (
+    AmazonStability3TextToImageRequest,
+    AmazonStability3TextToImageResponse,
+)
+from litellm.types.utils import ImageResponse
+
+
+class AmazonStability3Config:
+    """
+    Reference: https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/providers?model=stability.stable-diffusion-xl-v0
+
+    Stability API Ref: https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post
+    """
+
+    @classmethod
+    def get_config(cls):
+        return {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not k.startswith("__")
+            and not isinstance(
+                v,
+                (
+                    types.FunctionType,
+                    types.BuiltinFunctionType,
+                    classmethod,
+                    staticmethod,
+                ),
+            )
+            and v is not None
+        }
+
+    @classmethod
+    def get_supported_openai_params(cls, model: Optional[str] = None) -> List:
+        """
+        No additional OpenAI params are mapped for stability 3
+        """
+        return []
+
+    @classmethod
+    def _is_stability_3_model(cls, model: Optional[str] = None) -> bool:
+        """
+        Returns True if the model is a Stability 3 model
+
+        Stability 3 models follow this pattern:
+            sd3-large
+            sd3-large-turbo
+            sd3-medium
+            sd3.5-large
+            sd3.5-large-turbo
+
+        Stability ultra models
+            stable-image-ultra-v1
+        """
+        if model:
+            if "sd3" in model or "sd3.5" in model:
+                return True
+            if "stable-image-ultra-v1" in model:
+                return True
+        return False
+
+    @classmethod
+    def transform_request_body(
+        cls, prompt: str, optional_params: dict
+    ) -> AmazonStability3TextToImageRequest:
+        """
+        Transform the request body for the Stability 3 models
+        """
+        data = AmazonStability3TextToImageRequest(prompt=prompt, **optional_params)
+        return data
+
+    @classmethod
+    def map_openai_params(cls, non_default_params: dict, optional_params: dict) -> dict:
+        """
+        Map the OpenAI params to the Bedrock params
+
+        No OpenAI params are mapped for Stability 3, so directly return the optional_params
+        """
+        return optional_params
+
+    @classmethod
+    def transform_response_dict_to_openai_response(
+        cls, model_response: ImageResponse, response_dict: dict
+    ) -> ImageResponse:
+        """
+        Transform the response dict to the OpenAI response
+        """
+
+        stability_3_response = AmazonStability3TextToImageResponse(**response_dict)
+        openai_images: List[Image] = []
+        for _img in stability_3_response.get("images", []):
+            openai_images.append(Image(b64_json=_img))
+
+        model_response.data = openai_images
+        return model_response
diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/cost_calculator.py b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/cost_calculator.py
new file mode 100644
index 00000000..0a20b44c
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/cost_calculator.py
@@ -0,0 +1,41 @@
+from typing import Optional
+
+import litellm
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+    model: str,
+    image_response: ImageResponse,
+    size: Optional[str] = None,
+    optional_params: Optional[dict] = None,
+) -> float:
+    """
+    Bedrock image generation cost calculator
+
+    Handles both Stability 1 and Stability 3 models
+    """
+    if litellm.AmazonStability3Config()._is_stability_3_model(model=model):
+        pass
+    else:
+        # Stability 1 models
+        optional_params = optional_params or {}
+
+        # see model_prices_and_context_window.json for details on how steps is used
+        # Reference pricing by steps for stability 1: https://aws.amazon.com/bedrock/pricing/
+        _steps = optional_params.get("steps", 50)
+        steps = "max-steps" if _steps > 50 else "50-steps"
+
+        # size is stored in model_prices_and_context_window.json as 1024-x-1024
+        # current size has 1024x1024
+        size = size or "1024-x-1024"
+        model = f"{size}/{steps}/{model}"
+
+    _model_info = litellm.get_model_info(
+        model=model,
+        custom_llm_provider="bedrock",
+    )
+
+    output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
+    num_images: int = len(image_response.data)
+    return output_cost_per_image * num_images
diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/image_handler.py b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/image_handler.py
new file mode 100644
index 00000000..8f7762e5
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/image/image_handler.py
@@ -0,0 +1,314 @@
+import copy
+import json
+import os
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import httpx
+from pydantic import BaseModel
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
+from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
+    HTTPHandler,
+    _get_httpx_client,
+    get_async_httpx_client,
+)
+from litellm.types.utils import ImageResponse
+
+from ..base_aws_llm import BaseAWSLLM
+from ..common_utils import BedrockError
+
+if TYPE_CHECKING:
+    from botocore.awsrequest import AWSPreparedRequest
+else:
+    AWSPreparedRequest = Any
+
+
+class BedrockImagePreparedRequest(BaseModel):
+    """
+    Internal/Helper class for preparing the request for bedrock image generation
+    """
+
+    endpoint_url: str
+    prepped: AWSPreparedRequest
+    body: bytes
+    data: dict
+
+
+class BedrockImageGeneration(BaseAWSLLM):
+    """
+    Bedrock Image Generation handler
+    """
+
+    def image_generation(
+        self,
+        model: str,
+        prompt: str,
+        model_response: ImageResponse,
+        optional_params: dict,
+        logging_obj: LitellmLogging,
+        timeout: Optional[Union[float, httpx.Timeout]],
+        aimg_generation: bool = False,
+        api_base: Optional[str] = None,
+        extra_headers: Optional[dict] = None,
+        client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
+    ):
+        prepared_request = self._prepare_request(
+            model=model,
+            optional_params=optional_params,
+            api_base=api_base,
+            extra_headers=extra_headers,
+            logging_obj=logging_obj,
+            prompt=prompt,
+        )
+
+        if aimg_generation is True:
+            return self.async_image_generation(
+                prepared_request=prepared_request,
+                timeout=timeout,
+                model=model,
+                logging_obj=logging_obj,
+                prompt=prompt,
+                model_response=model_response,
+                client=(
+                    client
+                    if client is not None and isinstance(client, AsyncHTTPHandler)
+                    else None
+                ),
+            )
+
+        if client is None or not isinstance(client, HTTPHandler):
+            client = _get_httpx_client()
+        try:
+            response = client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body)  # type: ignore
+            response.raise_for_status()
+        except httpx.HTTPStatusError as err:
+            error_code = err.response.status_code
+            raise BedrockError(status_code=error_code, message=err.response.text)
+        except httpx.TimeoutException:
+            raise BedrockError(status_code=408, message="Timeout error occurred.")
+        ### FORMAT RESPONSE TO OPENAI FORMAT ###
+        model_response = self._transform_response_dict_to_openai_response(
+            model_response=model_response,
+            model=model,
+            logging_obj=logging_obj,
+            prompt=prompt,
+            response=response,
+            data=prepared_request.data,
+        )
+        return model_response
+
+    async def async_image_generation(
+        self,
+        prepared_request: BedrockImagePreparedRequest,
+        timeout: Optional[Union[float, httpx.Timeout]],
+        model: str,
+        logging_obj: LitellmLogging,
+        prompt: str,
+        model_response: ImageResponse,
+        client: Optional[AsyncHTTPHandler] = None,
+    ) -> ImageResponse:
+        """
+        Asynchronous handler for bedrock image generation
+
+        Awaits the response from the bedrock image generation endpoint
+        """
+        async_client = client or get_async_httpx_client(
+            llm_provider=litellm.LlmProviders.BEDROCK,
+            params={"timeout": timeout},
+        )
+
+        try:
+            response = await async_client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body)  # type: ignore
+            response.raise_for_status()
+        except httpx.HTTPStatusError as err:
+            error_code = err.response.status_code
+            raise BedrockError(status_code=error_code, message=err.response.text)
+        except httpx.TimeoutException:
+            raise BedrockError(status_code=408, message="Timeout error occurred.")
+
+        ### FORMAT RESPONSE TO OPENAI FORMAT ###
+        model_response = self._transform_response_dict_to_openai_response(
+            model=model,
+            logging_obj=logging_obj,
+            prompt=prompt,
+            response=response,
+            data=prepared_request.data,
+            model_response=model_response,
+        )
+        return model_response
+
+    def _prepare_request(
+        self,
+        model: str,
+        optional_params: dict,
+        api_base: Optional[str],
+        extra_headers: Optional[dict],
+        logging_obj: LitellmLogging,
+        prompt: str,
+    ) -> BedrockImagePreparedRequest:
+        """
+        Prepare the request body, headers, and endpoint URL for the Bedrock Image Generation API
+
+        Args:
+            model (str): The model to use for the image generation
+            optional_params (dict): The optional parameters for the image generation
+            api_base (Optional[str]): The base URL for the Bedrock API
+            extra_headers (Optional[dict]): The extra headers to include in the request
+            logging_obj (LitellmLogging): The logging object to use for logging
+            prompt (str): The prompt to use for the image generation
+        Returns:
+            BedrockImagePreparedRequest: The prepared request object
+
+        The BedrockImagePreparedRequest contains:
+            endpoint_url (str): The endpoint URL for the Bedrock Image Generation API
+            prepped (httpx.Request): The prepared request object
+            body (bytes): The request body
+        """
+        try:
+            from botocore.auth import SigV4Auth
+            from botocore.awsrequest import AWSRequest
+        except ImportError:
+            raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
+        boto3_credentials_info = self._get_boto_credentials_from_optional_params(
+            optional_params, model
+        )
+
+        ### SET RUNTIME ENDPOINT ###
+        modelId = model
+        _, proxy_endpoint_url = self.get_runtime_endpoint(
+            api_base=api_base,
+            aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint,
+            aws_region_name=boto3_credentials_info.aws_region_name,
+        )
+        proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke"
+        sigv4 = SigV4Auth(
+            boto3_credentials_info.credentials,
+            "bedrock",
+            boto3_credentials_info.aws_region_name,
+        )
+
+        data = self._get_request_body(
+            model=model, prompt=prompt, optional_params=optional_params
+        )
+
+        # Make POST Request
+        body = json.dumps(data).encode("utf-8")
+
+        headers = {"Content-Type": "application/json"}
+        if extra_headers is not None:
+            headers = {"Content-Type": "application/json", **extra_headers}
+        request = AWSRequest(
+            method="POST", url=proxy_endpoint_url, data=body, headers=headers
+        )
+        sigv4.add_auth(request)
+        if (
+            extra_headers is not None and "Authorization" in extra_headers
+        ):  # prevent sigv4 from overwriting the auth header
+            request.headers["Authorization"] = extra_headers["Authorization"]
+        prepped = request.prepare()
+
+        ## LOGGING
+        logging_obj.pre_call(
+            input=prompt,
+            api_key="",
+            additional_args={
+                "complete_input_dict": data,
+                "api_base": proxy_endpoint_url,
+                "headers": prepped.headers,
+            },
+        )
+        return BedrockImagePreparedRequest(
+            endpoint_url=proxy_endpoint_url,
+            prepped=prepped,
+            body=body,
+            data=data,
+        )
+
+    def _get_request_body(
+        self,
+        model: str,
+        prompt: str,
+        optional_params: dict,
+    ) -> dict:
+        """
+        Get the request body for the Bedrock Image Generation API
+
+        Checks the model/provider and transforms the request body accordingly
+
+        Returns:
+            dict: The request body to use for the Bedrock Image Generation API
+        """
+        provider = model.split(".")[0]
+        inference_params = copy.deepcopy(optional_params)
+        inference_params.pop(
+            "user", None
+        )  # make sure user is not passed in for bedrock call
+        data = {}
+        if provider == "stability":
+            if litellm.AmazonStability3Config._is_stability_3_model(model):
+                request_body = litellm.AmazonStability3Config.transform_request_body(
+                    prompt=prompt, optional_params=optional_params
+                )
+                return dict(request_body)
+            else:
+                prompt = prompt.replace(os.linesep, " ")
+                ## LOAD CONFIG
+                config = litellm.AmazonStabilityConfig.get_config()
+                for k, v in config.items():
+                    if (
+                        k not in inference_params
+                    ):  # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
+                        inference_params[k] = v
+                data = {
+                    "text_prompts": [{"text": prompt, "weight": 1}],
+                    **inference_params,
+                }
+        elif provider == "amazon":
+            return dict(litellm.AmazonNovaCanvasConfig.transform_request_body(text=prompt, optional_params=optional_params))
+        else:
+            raise BedrockError(
+                status_code=422, message=f"Unsupported model={model}, passed in"
+            )
+        return data
+
+    def _transform_response_dict_to_openai_response(
+        self,
+        model_response: ImageResponse,
+        model: str,
+        logging_obj: LitellmLogging,
+        prompt: str,
+        response: httpx.Response,
+        data: dict,
+    ) -> ImageResponse:
+        """
+        Transforms the Image Generation response from Bedrock to OpenAI format
+        """
+
+        ## LOGGING
+        if logging_obj is not None:
+            logging_obj.post_call(
+                input=prompt,
+                api_key="",
+                original_response=response.text,
+                additional_args={"complete_input_dict": data},
+            )
+        verbose_logger.debug("raw model_response: %s", response.text)
+        response_dict = response.json()
+        if response_dict is None:
+            raise ValueError("Error in response object format, got None")
+
+        config_class = (
+            litellm.AmazonStability3Config
+            if litellm.AmazonStability3Config._is_stability_3_model(model=model)
+            else litellm.AmazonNovaCanvasConfig if litellm.AmazonNovaCanvasConfig._is_nova_model(model=model)
+            else litellm.AmazonStabilityConfig
+        )
+        config_class.transform_response_dict_to_openai_response(
+            model_response=model_response,
+            response_dict=response_dict,
+        )
+
+        return model_response