google/generativeai/notebook/magics_engine.py

# -*- coding: utf-8 -*-
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""MagicsEngine class."""
from __future__ import annotations

from typing import AbstractSet, Sequence

from google.generativeai.notebook import argument_parser
from google.generativeai.notebook import cmd_line_parser
from google.generativeai.notebook import command
from google.generativeai.notebook import compare_cmd
from google.generativeai.notebook import compile_cmd
from google.generativeai.notebook import eval_cmd
from google.generativeai.notebook import ipython_env
from google.generativeai.notebook import model_registry
from google.generativeai.notebook import parsed_args_lib
from google.generativeai.notebook import post_process_utils
from google.generativeai.notebook import run_cmd
from google.generativeai.notebook.lib import prompt_utils


class MagicsEngine:
    """Implementation of functionality used by Magics.

    This class provides the implementation for Magics, decoupled from the
    details of integrating with Colab Magics such as registration.
    """

    def __init__(
        self,
        registry: model_registry.ModelRegistry | None = None,
        env: ipython_env.IPythonEnv | None = None,
    ):
        self._ipython_env = env
        models = registry or model_registry.ModelRegistry()
        self._cmd_handlers: dict[parsed_args_lib.CommandName, command.Command] = {
            parsed_args_lib.CommandName.RUN_CMD: run_cmd.RunCommand(models=models, env=env),
            parsed_args_lib.CommandName.COMPILE_CMD: compile_cmd.CompileCommand(
                models=models, env=env
            ),
            parsed_args_lib.CommandName.COMPARE_CMD: compare_cmd.CompareCommand(env=env),
            parsed_args_lib.CommandName.EVAL_CMD: eval_cmd.EvalCommand(models=models, env=env),
        }

    def parse_line(
        self,
        line: str,
        placeholders: AbstractSet[str],
    ) -> tuple[parsed_args_lib.ParsedArgs, parsed_args_lib.PostProcessingTokens]:
        return cmd_line_parser.CmdLineParser().parse_line(line, placeholders)

    def _get_handler(self, line: str, placeholders: AbstractSet[str]) -> tuple[
        command.Command,
        parsed_args_lib.ParsedArgs,
        Sequence[post_process_utils.ParsedPostProcessExpr],
    ]:
        """Given the command line, parse and return all components.

        Args:
          line: The LLM Magics command line.
          placeholders: Placeholders from prompts in the cell contents.

        Returns:
          A three-tuple containing:
          - The command (e.g. "run")
          - Parsed arguments for the command,
          - Parsed post-processing expressions
        """
        parsed_args, post_processing_tokens = self.parse_line(line, placeholders)
        cmd_name = parsed_args.cmd
        handler = self._cmd_handlers[cmd_name]
        post_processing_fns = handler.parse_post_processing_tokens(post_processing_tokens)
        return handler, parsed_args, post_processing_fns

    def execute_cell(self, line: str, cell_content: str):
        """Executes the supplied magic line and cell payload."""
        cell = _clean_cell(cell_content)
        placeholders = prompt_utils.get_placeholders(cell)

        try:
            handler, parsed_args, post_processing_fns = self._get_handler(line, placeholders)
            return handler.execute(parsed_args, cell, post_processing_fns)
        except argument_parser.ParserNormalExit as e:
            if self._ipython_env is not None:
                e.set_ipython_env(self._ipython_env)
            # ParserNormalExit implements the _ipython_display_ method so it can
            # be returned as the output of this cell for display.
            return e
        except argument_parser.ParserError as e:
            e.display(self._ipython_env)
            # Raise an exception to indicate that execution for this cell has
            # failed.
            # The exception is re-raised as SystemExit because Colab automatically
            # suppresses traceback for SystemExit but not other exceptions. Because
            # ParserErrors are usually due to user error (e.g. a missing required
            # flag or an invalid flag value), we want to hide the traceback to
            # avoid detracting the user from the error message, and we want to
            # reserve exceptions-with-traceback for actual bugs and unexpected
            # errors.
            error_msg = "Got parser error: {}".format(e.msgs()[-1]) if e.msgs() else ""
            raise SystemExit(error_msg) from e


def _clean_cell(cell_content: str) -> str:
    # Colab includes a trailing newline in cell_content. Remove only the last
    # line break from cell contents (i.e. not rstrip), so that multi-line and
    # intentional line breaks are preserved, but single-line prompts don't have
    # a trailing line break.
    cell = cell_content
    if cell.endswith("\n"):
        cell = cell[:-1]
    return cell