afnio.models

class ChatCompletionModel(BaseModel):
    """
    An abstraction for a language model that accepts a prompt composed of an array
    of messages containing instructions for the model. Each message can have a
    different role, influencing how the model interprets the input.
    """

    def __init__(self, provider: str = None, **kwargs):
        """Initializes the `ChatCompletionModel` instance.

        Args:
            provider: The name of the model provider (e.g., `"openai"`, `"anthropic"`).
            **kwargs: Recognized/expected keys are `usage` (a dictionary to track token
                usage and cost information) and any provider-specific configuration
                parameters.
        """
        super().__init__(provider=provider, **kwargs)

    # TODO: Add link to `API documentation` for kwargs of each supported model
    async def achat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """
        Asynchronous method to handle chat-based interactions with the model.

        Args:
            messages: A list of messages, where each message is represented as a
                dictionary with `"role"` (e.g., `"user"`, `"system"`) and `"content"`
                (the text of the message).
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A string containing the model's response to the chat messages.
        """
        raise NotImplementedError

    def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """
        Synchronous method to handle chat-based interactions with the model.

        Args:
            messages: A list of messages, where each message is represented as a
                dictionary with `"role"` (e.g., `"user"`, `"system"`) and `"content"`
                (the text of the message).
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A string containing the model's response to the chat messages.
        """
        raise NotImplementedError

class EmbeddingModel(BaseModel):
    """
    An abstraction for a model that generates embeddings for input texts.
    """

    def __init__(self, provider: str = None, **kwargs):
        """Initializes the `EmbeddingModel` instance.

        Args:
            provider: The name of the model provider (e.g., `"openai"`, `"anthropic"`).
            **kwargs: Recognized/expected keys are `usage` (a dictionary to track token
                usage and cost information) and any provider-specific configuration
                parameters.
        """
        super().__init__(provider=provider, **kwargs)

    async def aembed(self, input: List[str], **kwargs) -> List[List[float]]:
        """
        Asynchronous method to generate embeddings for the given input texts.

        Args:
            input: A list of input strings for which embeddings should be generated.
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A list of embeddings, where each embedding is represented \
            as a list of floats corresponding to the input strings.
        """
        raise NotImplementedError

    def embed(self, input: List[str], **kwargs) -> List[List[float]]:
        """
        Synchronous method to generate embeddings for the given input texts.

        Args:
            input: A list of input strings for which embeddings should be generated.
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A list of embeddings, where each embedding is represented \
            as a list of floats corresponding to the input strings.
        """
        raise NotImplementedError

class TextCompletionModel(BaseModel):
    """
    An abstraction for a language model that accepts a prompt composed of a single
    text input and generates a textual completion.
    """

    def __init__(self, provider: str = None, **kwargs):
        """Initializes the `TextCompletionModel` instance.

        Args:
            provider: The name of the model provider (e.g., `"openai"`, `"anthropic"`).
            **kwargs: Recognized/expected keys are `usage` (a dictionary to track token
                usage and cost information) and any provider-specific configuration
                parameters.
        """
        super().__init__(provider=provider, **kwargs)

    async def acomplete(self, prompt: str, **kwargs) -> str:
        """
        Asynchronous method to generate a completion for the given prompt.

        Args:
            prompt: The input text for which the model should generate a completion.
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A string containing the generated completion.
        """
        raise NotImplementedError

    def complete(self, prompt: str, **kwargs) -> str:
        """
        Synchronous method to generate a completion for the given prompt.

        Args:
            prompt: The input text for which the model should generate a completion.
            **kwargs: Additional parameters to configure the model's behavior during
                chat completion. This may include options such as:

                - model (`str`): The model to use (e.g., `"gpt-4o"`).
                - temperature (`float`): Amount of randomness injected into
                    the response.
                - max_completion_tokens (`int`): Maximum number of tokens to generate.
                - etc.

                For a complete list of supported parameters for each model, refer to the
                respective API documentation.

        Returns:
            A string containing the generated completion.
        """
        raise NotImplementedError