evalstone/evalscope/evalscope/api/model/model.py

import abc
import asyncio
from functools import partial
from pydantic_core import to_jsonable_python
from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Optional, Sequence, Union

from evalscope.api.messages import ChatMessage, ChatMessageAssistant, ChatMessageSystem, ChatMessageUser
from evalscope.api.registry import get_model_api
from evalscope.api.tool import ToolChoice, ToolFunction, ToolInfo
from evalscope.utils import get_logger, get_secret_value
from evalscope.utils.function_utils import thread_safe
from .generate_config import GenerateConfig
from .model_output import ModelOutput

if TYPE_CHECKING:
    from evalscope.config import TaskConfig

logger = get_logger()


class ModelAPI(abc.ABC):
    """Model API provider."""

    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        config: GenerateConfig = GenerateConfig(),
        **kwargs
    ) -> None:
        """Create a model API provider.

        Args:
           model_name (str): Model name.
           base_url (str | None): Alternate base URL for model.
           api_key (str | None): API key for model.
           api_key_vars (list[str]): Environment variables that
              may contain keys for this provider (used for override)
           config (GenerateConfig): Model configuration.
        """
        self.model_name = model_name
        self.base_url = base_url
        self.api_key = api_key
        self.config = config

    @abc.abstractmethod
    def generate(
        self,
        input: List[ChatMessage],
        tools: List[ToolInfo],
        tool_choice: ToolChoice,
        config: GenerateConfig,
    ) -> ModelOutput:
        """Generate output from the model.

        Args:
          input (str | list[ChatMessage]): Chat message
            input (if a `str` is passed it is converted
            to a `ChatUserMessage`).
          tools (list[ToolInfo]): Tools available for the
            model to call.
          tool_choice (ToolChoice): Directives to the model
            as to which tools to prefer.
          config (GenerateConfig): Model configuration.

        Returns:
           ModelOutput
        """
        ...

    async def generate_async(
        self,
        input: List[ChatMessage],
        tools: List[ToolInfo],
        tool_choice: ToolChoice,
        config: GenerateConfig,
    ) -> ModelOutput:
        """Async version of generate().

        Default implementation wraps the synchronous generate() in a thread
        pool executor so it does not block the event loop.  Subclasses (e.g.
        OpenAI-compatible models) may override this with a native async
        implementation for better performance.

        Args:
          input: Chat message input.
          tools: Tools available for the model to call.
          tool_choice: Directives to the model as to which tools to prefer.
          config: Model configuration.

        Returns:
           ModelOutput
        """
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None,
            partial(
                self.generate,
                input=input,
                tools=tools,
                tool_choice=tool_choice,
                config=config,
            ),
        )

    def max_tokens(self) -> Optional[int]:
        """Default max_tokens."""
        return None

    def max_tokens_for_config(self, config: GenerateConfig) -> Optional[int]:
        """Default max_tokens for a given config.

        Args:
           config: Generation config.

        Returns:
           Default maximum tokens for specified configuration.
        """
        return None

    def tools_required(self) -> bool:
        """Any tool use in a message stream means that tools must be passed."""
        return False

    def tool_result_images(self) -> bool:
        """Tool results can contain images"""
        return False


class Model:
    """Model interface.

    Use `get_model()` to get an instance of a model.
    """

    api: ModelAPI
    """Model API."""

    config: GenerateConfig
    """Generation config."""

    def __init__(self, api: ModelAPI, config: GenerateConfig, model_args: Dict[str, Any] = {}) -> None:
        """Create a model.

        Args:
           api: Model API provider.
           config: Model configuration.
           model_args: Optional model args
        """
        self.api = api
        self.config = config
        self.model_args = model_args

    @property
    def name(self) -> str:
        """Model name or path to model."""
        return self.api.model_name

    @property
    def role(self) -> Optional[str]:
        """Model role."""
        return self._role

    @role.setter
    def role(self, role: str) -> None:
        self._role = role

    def __str__(self) -> str:
        return f'Model(name={self.name}, role={self.role})'

    def generate(
        self,
        input: Union[str, List[ChatMessage]],
        tools: Optional[Sequence[ToolInfo]] = None,
        tool_choice: Optional[ToolChoice] = None,
        config: Optional[GenerateConfig] = None,
    ) -> ModelOutput:
        """Generate output from the model.

        Args:
          input: Chat message input (if a `str` is passed it is converted
            to a `ChatMessageUser`).
          tools: Tools available for the model to call.
          tool_choice: Directives to the model as to which tools to prefer.
          config: Model configuration.

        Returns:
           ModelOutput
        """
        processed_input, processed_tools, processed_tool_choice, processed_config = self._preprocess_input(
            input, tools, tool_choice, config
        )

        # Call the model's generate method
        output = self.api.generate(
            input=processed_input,
            tools=processed_tools,
            tool_choice=processed_tool_choice,
            config=processed_config,
        )

        # return output
        return output

    async def generate_async(
        self,
        input: Union[str, List[ChatMessage]],
        tools: Optional[Sequence[ToolInfo]] = None,
        tool_choice: Optional[ToolChoice] = None,
        config: Optional[GenerateConfig] = None,
    ) -> ModelOutput:
        """Async version of generate().

        Delegates to the underlying ModelAPI.generate_async() after
        preprocessing input.  The default ModelAPI implementation offloads
        the synchronous generate() to a thread pool; concrete subclasses
        (e.g. OpenAICompatibleAPI) may provide a native async implementation.

        Args:
          input: Chat message input (if a `str` is passed it is converted
            to a `ChatMessageUser`).
          tools: Tools available for the model to call.
          tool_choice: Directives to the model as to which tools to prefer.
          config: Model configuration.

        Returns:
           ModelOutput
        """
        processed_input, processed_tools, processed_tool_choice, processed_config = self._preprocess_input(
            input, tools, tool_choice, config
        )

        # Call the model's async generate method
        output = await self.api.generate_async(
            input=processed_input,
            tools=processed_tools,
            tool_choice=processed_tool_choice,
            config=processed_config,
        )

        return output

    def _preprocess_input(
        self,
        input: Union[str, List[ChatMessage]],
        tools: Optional[Sequence[ToolInfo]] = None,
        tool_choice: Optional[ToolChoice] = None,
        config: Optional[GenerateConfig] = None,
    ) -> tuple[List[ChatMessage], List[ToolInfo], ToolChoice, GenerateConfig]:
        """pre process input for generate."""

        # merge passed config
        if config is not None:
            config = self.config.merge(config)
        else:
            config = self.config.model_copy(deep=True)

        # provide max_tokens from the model api if required
        if config.max_tokens is None:
            config.max_tokens = self.api.max_tokens_for_config(config)
            if config.max_tokens is None:
                config.max_tokens = self.api.max_tokens()

        # normalize input to chat
        if isinstance(input, str):
            input = [ChatMessageUser(content=input)]

        # handle tools and tool_choice
        tool_choice = tool_choice if tool_choice is not None else 'auto'
        tools_info = list(tools) if tools is not None else []

        if isinstance(tool_choice, ToolFunction):
            tools_info = [tool for tool in tools_info if tool.name == tool_choice.name]

        if tool_choice == 'none' or len(tools_info) == 0:
            if not self.api.tools_required():
                tools_info = []
            tool_choice = 'none'

        return input, tools_info, tool_choice, config


class ModelCache:
    _models: Dict[str, 'Model'] = {}

    @classmethod
    def get(cls, key: str) -> Optional['Model']:
        return cls._models.get(key, None)

    @classmethod
    def set(cls, key: str, model: 'Model') -> None:
        cls._models[key] = model


def get_model_with_task_config(task_config: 'TaskConfig') -> Model:
    """Get an instance of a model with the specified task configuration.

    Args:
        task_config (TaskConfig): Task configuration.

    Returns:
        Model: An instance of the model.
    """
    model = task_config.model
    eval_type = task_config.eval_type
    base_url = task_config.api_url
    api_key = get_secret_value(task_config.api_key)
    config = task_config.generation_config
    model_args = task_config.model_args or {}

    return get_model(
        model=model, eval_type=eval_type, base_url=base_url, api_key=api_key, config=config, model_args=model_args
    )


@thread_safe
def get_model(
    model: Union[str, Model, ModelAPI],
    eval_type: str,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    config: GenerateConfig = GenerateConfig(),
    model_args: dict = {},
    role: Optional[str] = None,
    memoize: bool = True,
) -> Model:
    """Get an instance of a model.

    Calls to get_model() are memoized (i.e. a call with the same arguments
    will return an existing instance of the model rather than creating a
    new one). You can disable this with `memoize=False`.

    Args:
        task_config (TaskConfig): Task configuration.
        memoize (bool): Whether to memoize the model instance.

    Returns:
        Model instance.

    """

    # start with seeing if a model was passed
    if isinstance(model, Model):
        return model

    if isinstance(model, ModelAPI):
        return Model(model, config, model_args)

    # see if we can return a memoized model instance
    # (exclude mockllm since custom_outputs is an infinite generator)
    model_cache_key: str = ''
    if eval_type.startswith('mock_llm'):
        memoize = False
    if memoize:
        model_cache_key = (
            model + str(role) + config.model_dump_json(exclude_none=True) + str(base_url) + str(api_key)
            + str(to_jsonable_python(model_args, fallback=lambda _: None))
        )
        cached = ModelCache.get(model_cache_key)
        if cached is not None:
            return cached

    # TaskConfig defaults api_key to 'EMPTY'; treat it as unset so env vars can be used.
    if api_key == 'EMPTY':
        api_key = None

    logger.info(
        f'Creating model {model} with eval_type={eval_type} '
        f'base_url={base_url}, config={config.model_dump(exclude_none=True)}, model_args={model_args}'
    )

    # find a matching model type
    modelapi_type = get_model_api(eval_type)

    modelapi_instance = modelapi_type(
        model_name=model,
        base_url=base_url,
        api_key=api_key,
        config=config,
        **model_args,
    )
    m = Model(modelapi_instance, config, model_args)
    if role is not None:
        m.role = role
    if memoize:
        ModelCache.set(model_cache_key, m)
    return m