from typing import Any, Callable, Dict, TypeVar

from langchain import BasePromptTemplate, LLMChain
from langchain.chat_models.base import BaseChatModel
from langchain.schema import BaseOutputParser, OutputParserException
from openai.error import (
    AuthenticationError,
    InvalidRequestError,
    RateLimitError,
    ServiceUnavailableError,
)

from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.web.api.errors import OpenAIError

T = TypeVar("T")


def parse_with_handling(parser: BaseOutputParser[T], completion: str) -> T:
    try:
        return parser.parse(completion)
    except OutputParserException as e:
        raise OpenAIError(
            e, "There was an issue parsing the response from the AI model."
        )


async def openai_error_handler(
    func: Callable[..., Any], *args: Any, settings: ModelSettings, **kwargs: Any
) -> Any:
    try:
        return await func(*args, **kwargs)
    except ServiceUnavailableError as e:
        raise OpenAIError(
            e,
            "OpenAI is experiencing issues. Visit "
            "https://status.openai.com/ for more info.",
            should_log=not settings.custom_api_key,
        )
    except InvalidRequestError as e:
        if e.user_message.startswith("The model:"):
            raise OpenAIError(
                e,
                f"Your API key does not have access to your current model. Please use a different model.",
                should_log=not settings.custom_api_key,
            )
        raise OpenAIError(e, e.user_message)
    except AuthenticationError as e:
        raise OpenAIError(
            e,
            "Authentication error: Ensure a valid API key is being used.",
            should_log=not settings.custom_api_key,
        )
    except RateLimitError as e:
        if e.user_message.startswith("You exceeded your current quota"):
            raise OpenAIError(
                e,
                f"Your API key exceeded your current quota, please check your plan and billing details.",
                should_log=not settings.custom_api_key,
            )
        raise OpenAIError(e, e.user_message)
    except Exception as e:
        raise OpenAIError(
            e, "There was an unexpected issue getting a response from the AI model."
        )


async def call_model_with_handling(
    model: BaseChatModel,
    prompt: BasePromptTemplate,
    args: Dict[str, str],
    settings: ModelSettings,
    **kwargs: Any,
) -> str:
    chain = LLMChain(llm=model, prompt=prompt)
    return await openai_error_handler(chain.arun, args, settings=settings, **kwargs)

langchain.LLMChain
[((469, 481), 'typing.TypeVar', 'TypeVar', (['"""T"""'], {}), "('T')\n", (476, 481), False, 'from typing import Any, Callable, Dict, TypeVar\n'), ((2486, 2520), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'model', 'prompt': 'prompt'}), '(llm=model, prompt=prompt)\n', (2494, 2520), False, 'from langchain import BasePromptTemplate, LLMChain\n'), ((662, 738), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', '"""There was an issue parsing the response from the AI model."""'], {}), "(e, 'There was an issue parsing the response from the AI model.')\n", (673, 738), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((993, 1138), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', '"""OpenAI is experiencing issues. Visit https://status.openai.com/ for more info."""'], {'should_log': '(not settings.custom_api_key)'}), "(e,\n    'OpenAI is experiencing issues. Visit https://status.openai.com/ for more info.'\n    , should_log=not settings.custom_api_key)\n", (1004, 1138), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((1522, 1552), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', 'e.user_message'], {}), '(e, e.user_message)\n', (1533, 1552), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((1604, 1729), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', '"""Authentication error: Ensure a valid API key is being used."""'], {'should_log': '(not settings.custom_api_key)'}), "(e,\n    'Authentication error: Ensure a valid API key is being used.',\n    should_log=not settings.custom_api_key)\n", (1615, 1729), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((2114, 2144), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', 'e.user_message'], {}), '(e, e.user_message)\n', (2125, 2144), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((2186, 2275), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', '"""There was an unexpected issue getting a response from the AI model."""'], {}), "(e,\n    'There was an unexpected issue getting a response from the AI model.')\n", (2197, 2275), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((1299, 1453), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', 'f"""Your API key does not have access to your current model. Please use a different model."""'], {'should_log': '(not settings.custom_api_key)'}), "(e,\n    f'Your API key does not have access to your current model. Please use a different model.'\n    , should_log=not settings.custom_api_key)\n", (1310, 1453), False, 'from reworkd_platform.web.api.errors import OpenAIError\n'), ((1892, 2045), 'reworkd_platform.web.api.errors.OpenAIError', 'OpenAIError', (['e', 'f"""Your API key exceeded your current quota, please check your plan and billing details."""'], {'should_log': '(not settings.custom_api_key)'}), "(e,\n    f'Your API key exceeded your current quota, please check your plan and billing details.'\n    , should_log=not settings.custom_api_key)\n", (1903, 2045), False, 'from reworkd_platform.web.api.errors import OpenAIError\n')]