# to be robust to multiple response formats.
            if isinstance(parsed_response, dict):
                output_keys = ["data", "output", "outputs"]
                for key in output_keys:
                    if key in parsed_response:
                        output_item = parsed_response[key]
                        break
                else:
                    raise ValueError(
                        f"No valid key ({', '.join(output_keys)}) in response:"
                        f" {parsed_response}"
                    )
                if isinstance(output_item, list):
                    text = output_item[0]
                else:
                    text = output_item
            else:
                raise ValueError(f"Unexpected response type: {parsed_response}")
            # Older versions of the API include the input in the output response
            if text.startswith(prompt):
                text = text[len(prompt) :]
        except requests.exceptions.JSONDecodeError as e:
            raise ValueError(
                f"Error raised by inference API: {e}.\nResponse: {response.text}"
            )
        # TODO: replace when MosaicML supports custom stop tokens natively
        if stop is not None:
            text = enforce_stop_tokens(text, stop)
        return text