executor.py

# WriterAgent - AI Writing Assistant for LibreOffice
# Copyright (c) 2024 John Balis
# Copyright (c) 2026 KeithCu (modifications and relicensing)
# Copyright (c) 2026 LibreCalc AI Assistant (Calc integration features, originally MIT)
#
# Based on code from the LibrePythonista project (Apache 2.0)
# Source: https://github.com/Amourspirit/python-libre-pythonista-ext/blob/main/oxt/pythonpath/libre_pythonista_lib/doc/calc/doc/sheet/cell/code/py_module.py
# and: https://github.com/Amourspirit/python-libre-pythonista-ext/blob/main/tests/test_code/test_code_ast.py
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""In-process Python execution for Calc/Writer (LocalPythonExecutor sandbox).

Each tool call uses a fresh executor instance so variables do not leak across invocations.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from plugin.contrib.smolagents.local_python_executor import ExecutionTimeoutError, InterpreterError, LocalPythonExecutor
from plugin.framework.tool import ToolBaseDummy

if TYPE_CHECKING:
    from plugin.framework.tool import ToolContext

from plugin.framework.errors import WriterAgentException
from plugin.scripting.import_policy import format_inprocess_import_policy_for_prompt
from plugin.scripting.sandbox import CALC_AUTHORIZED_IMPORTS
from plugin.calc.bridge import CalcBridge
from plugin.calc.manipulator import CellManipulator
from plugin.calc.inspector import CellInspector

log = logging.getLogger("writeragent.calc.python.executor")

_SPREADSHEET_SERVICE = "com.sun.star.sheet.SpreadsheetDocument"
_CALC_HELPER_ERROR = "lp, get_range, set_range, and data_range only work in a Calc spreadsheet."
# json.dumps accepts these. Anything else (set, bytes, range, generators, custom objects) raises TypeError.
_JSON_RESULT_TYPES = (str, int, float, bool, list, dict)


def _is_spreadsheet(doc: Any) -> bool:
    """True when *doc* is a Calc model.

    An unconfigured MagicMock's ``supportsService`` returns another mock.
    Treat that as Calc so existing tool tests keep the helpers. A real
    Writer document returns False for the spreadsheet service.
    """
    supports = getattr(doc, "supportsService", None)
    if not callable(supports):
        return True
    try:
        result = supports(_SPREADSHEET_SERVICE)
    except Exception:
        log.debug("spreadsheet service check failed", exc_info=True)
        return False
    if type(result).__name__ == "MagicMock":
        return True
    return bool(result)


class PythonExecutor:
    """Runs Python in LO's embedded interpreter with document helpers (stdlib-only imports)."""

    doc_url: str
    executor: LocalPythonExecutor

    def __init__(self, doc_url: str) -> None:
        self.doc_url = doc_url
        self.executor = LocalPythonExecutor(additional_authorized_imports=list(CALC_AUTHORIZED_IMPORTS))

    def inject_helpers(self, bridge: CalcBridge, manipulator: CellManipulator, inspector: CellInspector, *, spreadsheet: bool = True) -> None:
        """Injects document interaction helpers into the environment."""

        def _require_calc(helper: str) -> None:
            if spreadsheet:
                return
            # Fail the call on a non-spreadsheet. These helpers used to be
            # injected on Writer too, and lp_helper caught the RuntimeError
            # and returned None, so the tool reported success.
            raise WriterAgentException(f"{helper}: {_CALC_HELPER_ERROR}", code="UNSUPPORTED_OPERATION")

        def lp_helper(addr: str) -> Any:
            """Read values from the spreadsheet. Generic for cells or ranges."""
            _require_calc("lp")
            try:
                if ":" in addr:
                    data = inspector.read_range(addr)
                    return [[cell["value"] for cell in row] for row in data]
                sheet = bridge.get_active_sheet()
                return manipulator.safe_get_cell_value(sheet, addr)
            except ExecutionTimeoutError:
                # Let ExecutionTimeoutError leave the cell. The sandbox arms
                # SIGALRM once (local_python_executor.timeout). Swallowing it
                # here leaves the alarm unset, and a loop of lp('A1') never
                # returns to the UI thread.
                raise
            except Exception:
                log.exception("lp_helper failed for address %s", addr)
                return None

        def set_range_helper(addr: str, data: Any) -> str | dict[str, Any]:
            """Write values back to the spreadsheet."""
            _require_calc("set_range")
            return manipulator.write_formula_range(addr, data)

        self.executor.send_variables({"lp": lp_helper, "Sheet": lp_helper, "get_range": lp_helper, "set_range": set_range_helper})

    def format_result(self, result: Any) -> Any:
        """Processes the result for storage and display."""
        if result is None or isinstance(result, _JSON_RESULT_TYPES):
            return result
        # Stringify set, bytes, range, and generators too. They have no
        # __dict__, and json.dumps then raises TypeError when the tool
        # result is serialized.
        try:
            return f"<Result: {str(result)}>"
        except Exception:
            return f"<Object: {type(result).__name__}>"

    def execute_with_return(self, code_snippet: str) -> Any:
        """Execute *code_snippet*; return the last expression value (REPL-style ``_`` is not kept for the next call)."""
        try:
            code_output = self.executor(code_snippet)
            result = code_output.output
            return self.format_result(result)
        except InterpreterError as e:
            log.exception("Sandbox execution error: \n%s", code_snippet)
            raise WriterAgentException(f"Execution Error: {str(e)}", code="PYTHON_EXECUTION_ERROR") from e
        except WriterAgentException:
            raise
        except Exception as e:
            log.exception("Error executing Python code: \n%s", code_snippet)
            raise WriterAgentException(f"Execution Error: {str(e)}", code="PYTHON_EXECUTION_ERROR") from e


class ExecutePythonScript(ToolBaseDummy):
    """Executes Python in LibreOffice's sandbox (no numpy); fresh environment every call."""

    name: str | None = "execute_python_script"
    intent: str | None = "analyze"
    description: str = format_inprocess_import_policy_for_prompt() + " The value of the last expression is returned. Each call starts with a clean environment."
    parameters: dict[str, Any] | None = {
        "type": "object",
        "properties": {
            "code": {"type": "string", "description": "The Python code to execute."},
            "data_range": {"type": "string", "description": "Optional cell range; values are injected as data (flat list for one row or column)."},
            "target_range": {"type": "string", "description": "Optional: Cell range (e.g. 'A1') to write the script result to."},
        },
        "required": ["code"],
    }
    uno_services: list[str] | None = ["com.sun.star.sheet.SpreadsheetDocument", "com.sun.star.text.TextDocument"]
    is_mutation: bool | None = True

    def execute(self, ctx: ToolContext, **kwargs: Any) -> dict[str, Any]:
        code = kwargs.get("code", "")
        data_range = kwargs.get("data_range")
        target_range = kwargs.get("target_range")

        doc_url = ctx.doc.getURL() or "untitled"
        spreadsheet = _is_spreadsheet(ctx.doc)
        if not spreadsheet and (data_range or target_range):
            raise WriterAgentException(_CALC_HELPER_ERROR, code="UNSUPPORTED_OPERATION")
        executor = PythonExecutor(doc_url)

        bridge = CalcBridge(ctx.doc)
        manipulator = CellManipulator(bridge)
        inspector = CellInspector(bridge)
        executor.inject_helpers(bridge, manipulator, inspector, spreadsheet=spreadsheet)

        # Inject data if data_range is provided
        if data_range:
            try:
                raw_data = inspector.read_range(data_range)
                py_data = [[cell["value"] for cell in row] for row in raw_data]
                if len(py_data) == 1:
                    py_data = py_data[0]
                elif all(len(row) == 1 for row in py_data):
                    py_data = [row[0] for row in py_data]
                executor.executor.send_variables({"data": py_data})
            except Exception:
                log.exception("Failed to inject data_range %s", data_range)

        result = executor.execute_with_return(code)

        # write_formula_range returns a summary string or an array-formula dict.
        write_status: str | dict[str, Any] = ""
        if target_range and result is not None:
            try:
                write_status = manipulator.write_formula_range(target_range, result)
            except Exception as e:
                write_status = f"Warning: Failed to write to {target_range}: {str(e)}"

        return {"status": "ok", "result": result, "write_status": write_status}