azure-ai-evaluation 1.0.1__py3-none-any.whl → 1.13.5__py3-none-any.whl

azure/ai/evaluation/_legacy/prompty/_utils.py

@@ -0,0 +1,663 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# cspell:ignore apng, retriable
+
+import copy
+import os
+import re
+import json
+import base64
+from dataclasses import dataclass, is_dataclass, fields
+from logging import Logger
+from pathlib import Path
+from typing import (
+    Any,
+    AsyncGenerator,
+    Dict,
+    Final,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from jinja2 import Template
+from openai import AsyncStream
+from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionUserMessageParam
+from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAIError
+
+from azure.ai.evaluation._constants import DefaultOpenEncoding
+from azure.ai.evaluation._legacy.prompty._exceptions import (
+    InvalidInputError,
+    JinjaTemplateError,
+    PromptyException,
+)
+
+from azure.ai.evaluation._legacy.prompty._yaml_utils import load_yaml
+
+
+# region: Resolving references
+
+
+@dataclass
+class PromptyModelConfiguration:
+    """
+    A dataclass that represents a model config of prompty.
+
+    :param api: Type of the LLM request, default value is chat.
+    :type api: str
+    :param configuration: Prompty model connection configuration
+    :type configuration: dict
+    :param parameters: Params of the LLM request.
+    :type parameters: dict
+    :param response: Return the complete response or the first choice, default value is first.
+    :type response: str
+    """
+
+    configuration: dict
+    parameters: Dict[str, Any]
+    response: str = "first"
+    model: Optional[str] = None
+    # _overflow: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        if not isinstance(self.configuration, dict):
+            raise PromptyException("The configuration of the model must be a dictionary.")
+
+        if not self.model:
+            self.model = self.configuration.get("azure_deployment", None) or self.configuration.get("model", None)
+
+
+T = TypeVar("T")
+
+
+def dataclass_from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
+    """Helper function to make creating dataclass instances from dictionaries easier.
+    Unlike using cls(**data), this function will ignore any keys in the dictionary that
+    are not fields in the dataclass. If the dataclass optionally contains an _overflow
+    field, any extra key/value paris will be placed in that field.
+
+    This does no type checking and inspects only the key names.
+
+    :param Type[T] cls: The dataclass type to create.
+    :param Dict[str, Any] data: The dictionary to create the dataclass instance from.
+    :return: The dataclass instance.
+    :rtype: T
+    """
+    if not is_dataclass(cls):
+        raise ValueError("This function only works with @dataclass Types")
+
+    fields_set: Set[str] = {f.name for f in fields(cls)}
+
+    params: Dict[str, Any] = {}
+    overflow: Dict[str, Any] = {}
+
+    for key, value in data.items():
+        if key in fields_set:
+            params[key] = value
+        else:
+            overflow[key] = value
+
+    if "_overflow" in fields_set:
+        params["_overflow"] = overflow
+
+    return cast(T, cls(**params))
+
+
+def resolve_references(origin: Mapping[str, Any], base_path: Optional[Path] = None) -> Dict[str, Any]:
+    """Resolve all reference in the object.
+
+    :param Mapping[str, Any] origin: The object to resolve.
+    :param Path|None base_path: The base path to resolve the file reference.
+    :return: The resolved object.
+    :rtype: Dict[str, Any]
+    """
+
+    def _resolve_references(origin: Any, base_path: Optional[Path] = None) -> Any:
+        if isinstance(origin, str):
+            return _resolve_reference(origin, base_path=base_path)
+        if isinstance(origin, list):
+            return [_resolve_references(item, base_path=base_path) for item in origin]
+        if isinstance(origin, dict):
+            return {key: _resolve_references(value, base_path=base_path) for key, value in origin.items()}
+        return origin
+
+    return {k: _resolve_references(v, base_path=base_path) for k, v in origin.items()}
+
+
+def _resolve_reference(reference: str, base_path: Optional[Path] = None) -> Union[str, dict]:
+    """
+    Resolve the reference, two types are supported, env, file.
+    When the string format is ${env:ENV_NAME}, the environment variable value will be returned.
+    When the string format is ${file:file_path}, return the loaded json object.
+
+    :param str reference: The reference string.
+    :param Path|None base_path: The base path to resolve the file reference.
+    :return: The resolved reference.
+    :rtype: str | dict
+    """
+    pattern = r"\$\{(\w+):(.*)\}"
+    match = re.match(pattern, reference)
+    if match:
+        reference_type, value = match.groups()
+        if reference_type == "env":
+            return os.environ.get(value, reference)
+
+        if reference_type == "file":
+            if not Path(value).is_absolute() and base_path:
+                path = Path(base_path) / value
+            else:
+                path = Path(value)
+
+            if not path.exists():
+                raise PromptyException(f"Cannot find the reference file {value}.")
+
+            with open(path, "r", encoding=DefaultOpenEncoding.READ) as f:
+                if path.suffix.lower() == ".json":
+                    return json.load(f)
+                if path.suffix.lower() in [".yml", ".yaml"]:
+                    return load_yaml(f)
+                return f.read()
+
+        # TODO ralphe: logging?
+        # logger.warning(f"Unknown reference type {reference_type}, return original value {reference}.")
+        return reference
+
+    return reference
+
+
+def update_dict_recursively(origin_dict: Mapping[str, Any], overwrite_dict: Mapping[str, Any]) -> Dict[str, Any]:
+    updated_dict: Dict[str, Any] = {}
+    for k, v in overwrite_dict.items():
+        if isinstance(v, dict):
+            updated_dict[k] = update_dict_recursively(origin_dict.get(k, {}), v)
+        else:
+            updated_dict[k] = v
+    for k, v in origin_dict.items():
+        if k not in updated_dict:
+            updated_dict[k] = v
+    return updated_dict
+
+
+# endregion
+
+
+# region: Jinja template rendering
+
+VALID_ROLES = ["system", "user", "assistant", "function"]
+"""Valid roles for the OpenAI Chat API"""
+
+PROMPTY_ROLE_SEPARATOR_PATTERN = re.compile(
+    r"(?i)^\s*#?\s*(" + "|".join(VALID_ROLES) + r")\s*:\s*\n", flags=re.MULTILINE
+)
+"""Pattern to match the role separator in a prompty template"""
+
+MARKDOWN_IMAGE_PATTERN = re.compile(r"(?P<match>!\[[^\]]*\]\(.*?(?=\"|\))\))", flags=re.MULTILINE)
+"""Pattern to match markdown syntax for embedding images such as ![alt text](url).
+   This uses a 'hack' where by naming the capture group, using re.split() will cause
+   the named capture group to appear in the list of split parts"""
+
+IMAGE_URL_PARSING_PATTERN = re.compile(
+    r"^!\[(?P<alt_text>[^\]]+)\]\((?P<link>(?P<scheme>[^:]+(?=:))?:?(?P<mime_type>[^;]+(?=;))?;?(?P<data>[^\)]*))\)$"
+)
+"""Pattern used to parse the image URL from the markdown syntax. This caputres the following groups:
+    - alt_text: The alt text for the image
+    - link: The full link
+    - scheme: The scheme used in the link (e.g. data, http, https)
+    - mime_type: The mime type of the image (only for data URLs)
+    - data: The data part of the URL (only for data URLs)
+"""
+
+DEFAULT_IMAGE_MIME_TYPE: Final[str] = "image/*"
+"""The mime type to use when we don't know the image type"""
+
+FILE_EXT_TO_MIME: Final[Mapping[str, str]] = {
+    ".apng": "image/apng",
+    ".avif": "image/avif",
+    ".bmp": "image/bmp",
+    ".gif": "image/gif",
+    ".heic": "image/heic",
+    ".heif": "image/heif",
+    ".ico": "image/vnd.microsoft.icon",
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".png": "image/png",
+    ".svg": "image/svg+xml",
+    ".tif": "image/tiff",
+    ".tiff": "image/tiff",
+    ".webp": "image/webp",
+}
+"""Mapping of file extensions to mime types"""
+
+
+def render_jinja_template(template_str: str, *, trim_blocks=True, keep_trailing_newline=True, **kwargs) -> str:
+    try:
+        template = Template(template_str, trim_blocks=trim_blocks, keep_trailing_newline=keep_trailing_newline)
+        return template.render(**kwargs)
+    except Exception as e:  # pylint: disable=broad-except
+        raise PromptyException(f"Failed to render jinja template - {type(e).__name__}: {str(e)}") from e
+
+
+def build_messages(
+    *, prompt: str, working_dir: Path, image_detail: str = "auto", **kwargs: Any
+) -> Sequence[Mapping[str, Any]]:
+    # keep_trailing_newline=True is to keep the last \n in the prompt to avoid converting "user:\t\n" to "user:".
+    chat_str = render_jinja_template(prompt, trim_blocks=True, keep_trailing_newline=True, **kwargs)
+    messages = _parse_chat(chat_str, working_dir, image_detail)
+    return messages
+
+
+def _parse_chat(chat_str: str, working_dir: Path, image_detail: str) -> Sequence[Mapping[str, Any]]:
+    # openai chat api only supports VALID_ROLES as role names.
+    # customer can add single # in front of role name for markdown highlight.
+    # and we still support role name without # prefix for backward compatibility.
+
+    chunks = re.split(PROMPTY_ROLE_SEPARATOR_PATTERN, chat_str)
+    chat_list: List[Dict[str, Any]] = []
+
+    for chunk in chunks:
+        last_message = chat_list[-1] if len(chat_list) > 0 else None
+
+        # =======================================================================================================
+        # NOTE: The Promptflow code supported tool calls but used eval() to parse them. This is an unacceptable
+        #       security risk. Since none of the current evaluators use tool calls, this functionality has been
+        #       removed.
+        # =======================================================================================================
+
+        # if is_tool_chunk(last_message):
+        #     parse_tools(last_message, chunk, hash2images, image_detail)
+        #     continue
+        # if last_message and "role" in last_message and last_message["role"] == "assistant":
+        #     parsed_result = _try_parse_tool_calls(chunk)
+        #     if parsed_result is not None:
+        #         last_message["tool_calls"] = parsed_result
+        #         continue
+
+        if (
+            last_message
+            and "role" in last_message  # pylint: disable=unsupported-membership-test
+            and "content" not in last_message  # pylint: disable=unsupported-membership-test
+            and "tool_calls" not in last_message  # pylint: disable=unsupported-membership-test
+        ):
+            parsed_result = _try_parse_name_and_content(chunk)
+            if parsed_result is None:
+                if last_message["role"] == "function":  # pylint: disable=unsubscriptable-object
+                    # "name" is required if the role is "function"
+                    raise JinjaTemplateError(
+                        "Failed to parse function role prompt. Please make sure the prompt follows the "
+                        "format: 'name:\\nfunction_name\\ncontent:\\nfunction_content'. "
+                        "'name' is required if role is function, and it should be the name of the function "
+                        "whose response is in the content. May contain a-z, A-Z, 0-9, and underscores, "
+                        "with a maximum length of 64 characters. See more details in "
+                        "https://platform.openai.com/docs/api-reference/chat/create#chat/create-name"
+                    )
+
+                # "name" is optional for other role types.
+                last_message["content"] = _to_content_str_or_list(  # pylint: disable=unsupported-assignment-operation
+                    chunk, working_dir, image_detail
+                )
+            else:
+                last_message["name"] = parsed_result[0]  # pylint: disable=unsupported-assignment-operation
+                last_message["content"] = _to_content_str_or_list(  # pylint: disable=unsupported-assignment-operation
+                    parsed_result[1], working_dir, image_detail
+                )
+        else:
+            if chunk.strip() == "":
+                continue
+            # Check if prompt follows chat api message format and has valid role.
+            # References: https://platform.openai.com/docs/api-reference/chat/create.
+            role = chunk.strip().lower()
+            _validate_role(role)
+            new_message = {"role": role}
+            chat_list.append(new_message)
+    return chat_list
+
+
+def _validate_role(role: str):
+    if role not in VALID_ROLES:
+        valid_roles_str = ", ".join(VALID_ROLES)
+        error_message = (
+            f"The Chat API requires a specific format for prompt definition, and the prompt should include separate "
+            f"lines as role delimiters: {valid_roles_str}.\n"
+            f"Current parsed role '{role}' does not meet the requirement. If you intend to use the Completion API, "
+            f"please select the appropriate API type and deployment name."
+        )
+        raise JinjaTemplateError(message=error_message)
+
+
+def _to_content_str_or_list(chat_str: str, working_dir: Path, image_detail: str) -> Union[str, List[Dict[str, Any]]]:
+    chunks = [c for c in (chunk.strip() for chunk in re.split(MARKDOWN_IMAGE_PATTERN, chat_str)) if c]
+    if len(chunks) <= 1:
+        return chat_str.strip()
+
+    messages: List[Dict[str, Any]] = []
+    for chunk in chunks:
+        if chunk.startswith("![") and chunk.endswith(")"):
+            messages.append(_inline_image(chunk, working_dir, image_detail))
+        else:
+            messages.append({"type": "text", "text": chunk})
+    return messages
+
+
+def _inline_image(image: str, working_dir: Path, image_detail: str) -> Dict[str, Any]:
+    """This accepts an image URL in markdown format, and parses that into a message containing the image details
+    to be sent to AI service. In the case of local file images, they will be loaded and their contents encoded
+    into a base 64 data URI. Internal URLs will remained untouched. It can can accept http(s), ftp(s), as well
+    as data URIs.
+
+    :param str image: The image URL in markdown format (e.g. ![alternative text](https://www.bing.com/favicon.ico))
+    :param Path working_dir: The working directory to use when resolving relative file paths
+    :param str image_detail: The image detail to use when sending the image to the AI service
+    :return: The image message to send to the AI service
+    :rtype: Mapping[str, Any]"""
+
+    def local_to_base64(local_file: str, mime_type: Optional[str]) -> str:
+        path = Path(local_file)
+        if not path.is_absolute():
+            path = working_dir / local_file
+        if not path.exists():
+            # TODO ralphe logging?
+            # logger.warning(f"Cannot find the image path {image_content},
+            #                  it will be regarded as {type(image_str)}.")
+            raise InvalidInputError(f"Cannot find the image path '{path.as_posix()}'")
+
+        base64_encoded = base64.b64encode(path.read_bytes()).decode("utf-8")
+        if not mime_type:
+            mime_type = FILE_EXT_TO_MIME.get(path.suffix.lower(), DEFAULT_IMAGE_MIME_TYPE)
+        return f"data:{mime_type};base64,{base64_encoded}"
+
+    match = re.match(IMAGE_URL_PARSING_PATTERN, image)
+    if not match:
+        raise InvalidInputError(f"Invalid image URL '{image}'")
+
+    inlined_uri: str
+    mime_type: Optional[str] = None
+
+    scheme: str = (match.group("scheme") or "").strip().lower()
+    if scheme in ["http", "https", "ftp", "ftps"]:
+        # nothing special to do here, pass through full URI as is
+        inlined_uri = (match.group("link") or "").strip()
+    elif scheme == "data":
+        mime_type = (match.group("mime_type") or "").strip()
+        data: str = (match.group("data") or "").strip()
+
+        # data urls may contain local paths too
+        if data[:5].lower() == "path:":
+            inlined_uri = local_to_base64(data[5:].strip(), mime_type)
+        elif data[:6].lower() == "base64":
+            # nothing special to do here, pass through full URI as is
+            inlined_uri = (match.group("link") or "").strip()
+        else:
+            raise InvalidInputError(f"Invalid image data URL '{image}'")
+    else:
+        # assume it's a file path
+        inlined_uri = local_to_base64((match.group("link") or "").strip(), mime_type)
+
+    if not inlined_uri:
+        raise InvalidInputError(f"Failed to determine how to inline the following image URL '{image}'")
+
+    return {
+        "type": "image_url",
+        "image_url": {
+            "url": inlined_uri,
+            "detail": image_detail,
+        },
+    }
+
+
+def _try_parse_name_and_content(role_prompt: str) -> Optional[Tuple[str, str]]:
+    # customer can add ## in front of name/content for markdown highlight.
+    # and we still support name/content without ## prefix for backward compatibility.
+    # TODO ralphe: This maybe has something to do with parsing functions or tool calls but I'm not sure
+    pattern = r"\n*#{0,2}\s*name\s*:\s*\n+\s*(\S+)\s*\n*#{0,2}\s*content\s*:\s*\n?(.*)"
+    match = re.search(pattern, role_prompt, re.DOTALL)
+    if match:
+        return match.group(1), match.group(2)
+    return None
+
+
+# endregion
+
+
+# region OpenAI connections and requests
+
+OpenAIChatResponseType = Union[ChatCompletion, AsyncStream[ChatCompletionChunk]]
+
+
+def prepare_open_ai_request_params(
+    model_config: PromptyModelConfiguration, template: Union[str, Sequence[Mapping[str, Any]]]
+) -> MutableMapping[str, Any]:
+    params = copy.deepcopy(model_config.parameters)
+    # if isinstance(connection, AzureOpenAIConnection):
+    #     params.setdefault("extra_headers", {}).update({"ms-azure-ai-promptflow-called-from": "promptflow-core"})
+    params["model"] = model_config.model
+    params["messages"] = template
+
+    # NOTE:
+    # - Tool calls have been disabled due to a security issue in the implementation. See comment earlier in
+    #   this file for more details
+    # - Removing the validation of function calls in favour of letting the service do that validation. This
+    #   removes a maintenance burden from the SDK should the service definition for function calls change.
+
+    # # functions and function_call are deprecated and are replaced by tools and tool_choice.
+    # # if both are provided, tools and tool_choice are used and functions and function_call are ignored.
+    # if "tools" in params:
+    #     validate_tools(params["tools"])
+    #     params["tool_choice"] = validate_tool_choice(params.get("tool_choice", None))
+    # else:
+    #     if "functions" in params:
+    #         _validate_functions(params["functions"])
+    #         params["function_call"] = validate_function_call(params.get("function_call", None))
+
+    return params
+
+
+async def format_llm_response(
+    response: OpenAIChatResponseType,
+    is_first_choice: bool,
+    response_format: Optional[Mapping[str, Any]] = None,
+    outputs: Optional[Mapping[str, Any]] = None,
+    inputs: Optional[Mapping[str, Any]] = None,
+) -> dict:
+    """
+    Format LLM response
+
+    If is_first_choice is false, it will directly return LLM response.
+    If is_first_choice is true, behavior as blow:
+        response_format: type: text
+            - n: None/1/2
+                Return the first choice content. Return type is string.
+            - stream: True
+                Return generator list of first choice content. Return type is generator[str]
+        response_format: type: json_object
+            - n : None/1/2
+                Return json dict of the first choice. Return type is dict
+            - stream: True
+                Return json dict of the first choice. Return type is dict
+            - outputs
+                Extract corresponding output in the json dict to the first choice. Return type is dict.
+
+    :param response: LLM response.
+    :type response:
+    :param is_first_choice: If true, it will return the first item in response choices, else it will return all response
+    :type is_first_choice: bool
+    :param response_format: An object specifying the format that the model must output.
+    :type response_format: str
+    :param outputs: Extract corresponding output in json format response
+    :type outputs: dict
+    :return: Formatted LLM response.
+    :rtype: Union[str, dict, Response]
+    """
+
+    def format_choice(item: str) -> Union[str, Mapping[str, Any]]:
+        # response_format is one of text or json_object.
+        # https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format
+        if not is_json_format:
+            return item
+
+        result_dict = json.loads(item)
+        if not outputs:
+            return result_dict
+
+        # return the keys in outputs
+        output_results = {}
+        for key in outputs:
+            if key not in result_dict:
+                raise InvalidInputError(f"Cannot find '{key}' in response {list(result_dict.keys())}")
+            output_results[key] = result_dict[key]
+        return output_results
+
+    async def format_stream(llm_response: AsyncStream[ChatCompletionChunk]) -> AsyncGenerator[str, None]:
+        cur_index = None
+        async for chunk in llm_response:
+            if len(chunk.choices) > 0 and chunk.choices[0].delta.content:
+                if cur_index is None:
+                    cur_index = chunk.choices[0].index
+                if cur_index != chunk.choices[0].index:
+                    return
+                yield chunk.choices[0].delta.content
+
+    to_ret = {
+        "llm_output": None,
+        "input_token_count": 0,
+        "output_token_count": 0,
+        "total_token_count": 0,
+        "finish_reason": "",
+        "model_id": "",
+        "sample_input": "",
+        "sample_output": "",
+    }
+
+    if not is_first_choice:
+        to_ret["llm_output"] = response
+        return to_ret  # we don't actually use this code path since streaming is not used, so set token counts to 0
+
+    is_json_format = isinstance(response_format, dict) and response_format.get("type") == "json_object"
+    if isinstance(response, AsyncStream):
+        if not is_json_format:
+            to_ret["llm_output"] = format_stream(llm_response=response)
+            return to_ret
+        content = "".join([item async for item in format_stream(llm_response=response)])
+        to_ret["llm_output"] = format_choice(content)
+        return to_ret  # we don't actually use this code path since streaming is not used, so set token counts to 0
+    else:
+        input_token_count = response.usage.prompt_tokens if response.usage and response.usage.prompt_tokens else 0
+        output_token_count = (
+            response.usage.completion_tokens if response.usage and response.usage.completion_tokens else 0
+        )
+        total_token_count = response.usage.total_tokens if response.usage and response.usage.total_tokens else 0
+        finish_reason = (
+            response.choices[0].finish_reason if response.choices and response.choices[0].finish_reason else ""
+        )
+        model_id = response.model if response.model else ""
+        sample_output_list = (
+            [{"role": response.choices[0].message.role, "content": response.choices[0].message.content}]
+            if (response.choices and response.choices[0].message.content and response.choices[0].message.role)
+            else []
+        )
+        sample_output = json.dumps(sample_output_list)
+        input_str = f"{json.dumps(inputs)}" if inputs else ""
+        if inputs and len(inputs) > 0:
+            sample_input_json = []
+            msg = ChatCompletionUserMessageParam(
+                role="user",
+                content=input_str,
+            )
+            sample_input_json.append(msg)
+            sample_input = json.dumps(sample_input_json)
+
+    # When calling function/tool, function_call/tool_call response will be returned as a field in message,
+    # so we need return message directly. Otherwise, we only return content.
+    # https://platform.openai.com/docs/api-reference/chat/object#chat/object-choices
+    if response.choices[0].finish_reason in ["tool_calls", "function_calls"]:
+        response_content = response.model_dump()["choices"][0]["message"]
+    else:
+        response_content = getattr(response.choices[0].message, "content", "")
+    result = format_choice(response_content)
+    to_ret["llm_output"] = result
+    to_ret["input_token_count"] = input_token_count
+    to_ret["output_token_count"] = output_token_count
+    to_ret["total_token_count"] = total_token_count
+    to_ret["finish_reason"] = finish_reason
+    to_ret["model_id"] = model_id
+    to_ret["sample_input"] = sample_input
+    to_ret["sample_output"] = sample_output
+    return to_ret
+
+
+def openai_error_retryable(
+    error: OpenAIError, retry: int, entity_retry: List[int], max_entity_retries: int
+) -> Tuple[bool, float]:
+    """
+    Determines if an OpenAI error is retryable, and optionally determines the min retry delay to use.
+    If none is returned, the caller will determine the delay to use.
+
+    :param OpenAIError error: The error to handle
+    :param int retry: The current retry count (0 means we're on the first attempt and no retries have been made)
+    :param List[int] entity_retry: The current retry count for the unprocessable entity failures. This should be a
+        list containing only 1 element to mimic pass by reference semantics. A value of 0 means we're on the
+        first attempt and no retries have been made.
+    :param int max_entity_retries: The maximum number of retries to make for unprocessable entity failures
+    :return: A tuple containing whether the error is retryable and the min delay to use if any
+    :rtype: Tuple[bool, Optional[float]]
+    """
+
+    # Using https://platform.openai.com/docs/guides/error-codes/api-errors#python-library-error-types as a reference
+
+    should_retry: bool
+    delay: Optional[float] = None
+
+    if isinstance(error, APIConnectionError):
+        retriable_error_messages: Sequence[str] = [
+            "connection aborted",
+            # issue 2296
+            "server disconnected without sending a response",
+        ]
+        should_retry = (
+            isinstance(error, APITimeoutError)  # APITimeoutError is a subclass of APIConnectionError
+            or str(error).lower() in retriable_error_messages
+            or str(error.__cause__).lower() in retriable_error_messages
+        )
+    elif isinstance(error, APIStatusError):
+        status_code: int = error.response.status_code
+        if status_code == 422:
+            # As per the original legacy code, UnprocessableEntityError (HTTP 422) should be handled differently
+            # with a smaller retry count, as retrying more may not be beneficial.
+            should_retry = entity_retry[0] < max_entity_retries
+            entity_retry[0] += 1
+        elif status_code == 429:
+            # Two types, one is you are throttled and so should retry after a delay, the other is you have exceeded
+            # your quota and should not retry.
+            if (error.type or "").lower() == "insufficient_quota":
+                should_retry = False
+            else:
+                should_retry = True
+            should_retry = error.type != "insufficient_quota"
+        else:
+            should_retry = status_code >= 500
+
+        # Use what the service tells us to use for the delay if it's provided
+        if should_retry and not delay:
+            delay_str = error.response.headers.get("Retry-After", None)
+            if delay_str is not None:
+                delay = float(delay_str)
+    else:
+        should_retry = False
+
+    # Use exponential backoff for retries if the service doesn't provide a delay
+    if not delay:
+        delay = min(60, 2 + 2**retry)
+
+    return (should_retry, delay)
+
+
+# endregion