from .message_types import Msg
from .wrappers.llm_response_wrapper import LLMResponse
from ._env import env


def llm(prompt: str | Msg | list[str] | list[Msg], **kwargs) -> str | LLMResponse:
    """
    Request Large Language Model synchronously

    Args:
        prompt (str | list[str]): Text to send to LLM
        **kwargs (dict): Parameters supported by the LLM API

            See parameters supported by the OpenAI:

            - https://platform.openai.com/docs/api-reference/completions/create
            - https://platform.openai.com/docs/api-reference/chat/create

            **Additional parameters:**

                - callback: callable - callback function
                to be called on each chunk of text,
                enables response streaming if supported by the LLM API
                - callbacks: list[callable] - collection of callbacks
                to be called on each chunk of text,
                enables response streaming if supported by the LLM API

    Returns:

        Text generated by the LLM as string
        with all fields returned by API accessible as an attributes.

        See fields returned by the OpenAI:

        - https://platform.openai.com/docs/api-reference/completions/object
        - https://platform.openai.com/docs/api-reference/chat/object
    """
    [h(prompt, **kwargs) for h in env().llm_before_handlers]
    response = env().llm_function(prompt, **kwargs)
    [h(response) for h in env().llm_after_handlers]
    return response


async def allm(
    prompt: str | Msg | list[str] | list[Msg], **kwargs
) -> str | LLMResponse:
    """
    Request Large Language Model asynchronously

    Args:
        prompt (str | list[str]): Text to send to LLM
        **kwargs (dict): Parameters supported by the LLM API

            See parameters supported by the OpenAI:

            - https://platform.openai.com/docs/api-reference/completions/create
            - https://platform.openai.com/docs/api-reference/chat/create

            **Additional parameters:**

            - callback: callable - callback function
            to be called on each chunk of text,
            enables response streaming if supported by the LLM API
            - callbacks: list[callable] - collection of callbacks
            to be called on each chunk of text,
            enables response streaming if supported by the LLM API

            Note: async callbacks are supported only for async LLM API calls

    Returns:

        Text generated by the LLM as string
        with all fields returned by API accessible as an attributes.

        See fields returned by the OpenAI:

        - https://platform.openai.com/docs/api-reference/completions/object
        - https://platform.openai.com/docs/api-reference/chat/object
    """
    [h(prompt, **kwargs) for h in env().llm_before_handlers]
    response = await env().llm_async_function(prompt, **kwargs)
    [h(response) for h in env().llm_after_handlers]
    return response
