# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

from abc import ABC, abstractmethod
from typing import Any

Dialog = list[dict[str, Any]]


class IModel(ABC):
    """Base class for all LLM model wrappers."""

    def __init__(self, model_name: str, api_key: str | None = None):
        self.model_name = model_name
        self.api_key = api_key

    @abstractmethod
    def minibatch_size(self) -> int:
        raise NotImplementedError("Minibatch size is not implemented")

    @abstractmethod
    def generate(
        self,
        dialog_minibatch: list[Dialog],
        *,
        temperature: float,
        max_tokens: int,
        top_p: float | None,
    ) -> list[str]:
        """Generate a response from the model."""
        pass

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.model_name})"
