table.py

"""Tool for building and editing a literature review table.

Intended to support the arxivdigestables task.

Table assumes rows correspond to papers, as in the arxivdigestables task.

Detects and removes orphaned cells when papers are removed.

Highlights missing cells in the table after each operation.

"""

from typing import Annotated, Any, Literal

from inspect_ai.tool import Tool, ToolDef, ToolError, ToolResult, tool
from inspect_ai.util import StoreModel, store_as
from pydantic import BaseModel, Field, RootModel
from shortuuid import uuid

from .submission import get_submission_manager


class Cell(BaseModel):
    paper_id: str = Field(..., description="Paper ID for the cell")
    value: str = Field(..., description="Cell value")


class Column(BaseModel):
    title: str = Field(..., description="Column title")
    description: str = Field(..., description="Column description")
    values: list[Cell] = Field(
        default_factory=list,
        description="List of cells, each with paper_id and value",
    )


class TableData(BaseModel):
    columns: list[Column] = Field(default_factory=list, description="List of columns")
    paper_ids: list[str] = Field(
        default_factory=list, description="List of paper_ids in the table"
    )
    missing_cells: dict[str, list[str]] = Field(
        default_factory=dict,
        description="Mapping from column title to list of paper_ids missing a value",
    )


class AddPapersParams(BaseModel):
    action: Literal["add_papers"] = "add_papers"
    paper_ids: list[str] = Field(..., description="List of paper_id strings to append")


class RemovePapersParams(BaseModel):
    action: Literal["remove_papers"] = "remove_papers"
    paper_ids: list[str] = Field(..., description="List of paper_id strings to remove")


class AddColumnsParams(BaseModel):
    action: Literal["add_columns"] = "add_columns"
    columns: list[Column] = Field(
        ...,
        description="List of new columns to append; each may include title, description, and optional initial 'values' dict mapping paper_id→cell.",
    )
    before_column: str | None = Field(
        None,
        description=(
            "Title of the column before which to insert new columns; "
            "if omitted, appends to end"
        ),
    )


class EditColumnValuesParams(BaseModel):
    action: Literal["edit_column_values"] = "edit_column_values"
    column: str = Field(..., description="Title of the column to update")
    updates: list[Cell] = Field(
        ..., description="List of cells with paper_id and new value"
    )


class RemoveColumnsParams(BaseModel):
    action: Literal["remove_columns"] = "remove_columns"
    target_columns: list[str] = Field(
        ..., description="List of column titles to remove"
    )


class GetTableParams(BaseModel):
    action: Literal["get"] = "get"


class UndoParams(BaseModel):
    action: Literal["undo"] = "undo"


TableEditorParams = RootModel[
    Annotated[
        AddPapersParams
        | RemovePapersParams
        | AddColumnsParams
        | EditColumnValuesParams
        | RemoveColumnsParams
        | GetTableParams
        | UndoParams,
        Field(discriminator="action"),
    ]
]


class TableState(StoreModel):
    table_state: TableData = Field(default_factory=lambda: TableData())
    history: list[str] = Field(default_factory=list)


class TableResponsePayload(BaseModel):
    table: TableData
    orphaned_cells: list[Cell] = Field(default_factory=list)


@tool
def table_editor(
    instance: str | None = uuid(),
    history_limit: int = 3,
    enable_submission: bool = False,
) -> Tool:
    """Stateful tool to build and edit a table of papers.

    Args:
        enable_submission: controls whether the "submit" action is available.
            Submissions are available via `astabench.tools.submission` with
            `has_submission` and `get_submission`.
    """

    async def execute(
        action: Literal[
            "add_papers",
            "remove_papers",
            "add_columns",
            "edit_column_values",
            "remove_columns",
            "get",
            "undo",
            "submit",
        ],
        columns: list[Column] | None = None,
        paper_ids: list[str] | None = None,
        before_column: str | None = None,
        remove_columns: list[str] | None = None,
        edit_params: EditColumnValuesParams | None = None,
    ) -> ToolResult:
        """Execute table editing commands.

        Args:
          action: One of:
            - "add_papers": append new rows of paper_ids
            - "remove_papers": remove rows of paper_ids
            - "add_columns": add new columns
            - "edit_column_values": update the values of cells in a column
            - "remove_columns": remove columns by title
            - "get": retrieve current table state
            - "undo": revert to previous state
            - "submit": submit the table as the final answer
          columns: For "add_columns", list of Column objects. Can be just title/descriptions for outlining
          paper_ids: For "add_papers" and "remove_papers", list of paper_id strings
          before_column: For "add_columns", title of the column before which to insert new columns (default to end)
          remove_columns: For "remove_columns", list of column titles to delete
          edit_params: Params for "edit_column_values"

        Returns:
          A JSON-serializable dict with the current table state and any orphaned cells.

        """
        state = store_as(TableState, instance=instance)

        # prepare full rollback snapshot for mutating actions
        mutating = action not in ("get", "undo")
        if mutating:
            prior_history = state.history.copy()
            prior_table_json = state.table_state.model_dump_json()
            state.history.append(prior_table_json)
            # enforce maximum history length
            if len(state.history) > history_limit:
                state.history.pop(0)

        # Use a generic variable for params to satisfy mypy across cases
        p: Any

        # Track orphaned cells removed when papers are deleted
        orphaned_cells: list[Cell] = []

        # Validate params
        if action == "add_columns":
            if columns is None:
                raise ToolError("Missing columns for add_columns")
            p = AddColumnsParams(columns=columns, before_column=before_column)
        elif action == "add_papers":
            if paper_ids is None:
                raise ToolError("Missing paper_ids for add_papers")
            p = AddPapersParams(paper_ids=paper_ids)
        elif action == "edit_column_values":
            if edit_params is None:
                raise ToolError("Missing edit_params for edit_column_values")
            p = edit_params
        elif action == "remove_columns":
            if remove_columns is None:
                raise ToolError("Missing remove_columns for remove_columns")
            p = RemoveColumnsParams(target_columns=remove_columns)
        elif action == "remove_papers":
            if paper_ids is None:
                raise ToolError("Missing paper_ids for remove_papers")
            p = RemovePapersParams(paper_ids=paper_ids)
        elif action == "undo":
            p = UndoParams()
        elif action == "get":
            p = GetTableParams()
        elif action == "submit":
            if not enable_submission:
                raise ValueError("`submit` action not enabled")
            if (
                columns is not None
                or paper_ids is not None
                or before_column is not None
                or remove_columns is not None
                or edit_params is not None
            ):
                raise ToolError(
                    "Cannot submit table with pending changes. Please finalize edits and then use 'submit' with no other args."
                )
            if not state.table_state.columns or not state.table_state.paper_ids:
                raise ToolError("Cannot submit empty table")
            submission = state.table_state.model_dump_json()
            get_submission_manager().write_submission(submission)
            return TableResponsePayload(
                table=state.table_state,
                orphaned_cells=orphaned_cells,
            ).model_dump_json()
        else:
            raise ToolError(f"Unknown action {action!r}")

        # Dispatch with automatic rollback on any exception
        try:
            match p:
                case AddColumnsParams():
                    new_cols = []
                    for col in p.columns:
                        new_cols.append(col)
                    # ensure no duplicate titles in the bulk-add payload
                    titles = [col.title for col in new_cols]
                    if len(titles) != len(set(titles)):
                        raise ToolError("Duplicate titles in add_columns payload")
                    # prevent duplicate column titles
                    existing_titles = {c.title for c in state.table_state.columns}
                    for col in new_cols:
                        if col.title in existing_titles:
                            raise ToolError(
                                f"Duplicate column title: {col.title}. If you wanted to edit this column, use the edit_column action instead."
                            )
                        existing_titles.add(col.title)
                    if p.before_column is not None:
                        idx = next(
                            (
                                i
                                for i, c in enumerate(state.table_state.columns)
                                if c.title == p.before_column
                            ),
                            None,
                        )
                        if idx is None:
                            raise ToolError(
                                f"Column to insert before not found: {p.before_column}"
                            )
                        for offset, col in enumerate(new_cols):
                            state.table_state.columns.insert(idx + offset, col)
                    else:
                        state.table_state.columns.extend(new_cols)
                case AddPapersParams():
                    for pid in p.paper_ids:
                        if pid not in state.table_state.paper_ids:
                            state.table_state.paper_ids.append(pid)
                case EditColumnValuesParams():
                    # Update the specified column
                    for col in state.table_state.columns:
                        if col.title == p.column:
                            for cell_update in p.updates:
                                pid, new_val = cell_update.paper_id, cell_update.value
                                for cell in col.values:
                                    if cell.paper_id == pid:
                                        cell.value = new_val
                                        break
                                else:
                                    col.values.append(cell_update)
                            break
                    else:
                        raise ToolError(f"Column '{p.column}' not found")
                case RemoveColumnsParams():
                    # Drop columns from definition and from each column's values
                    state.table_state.columns = [
                        c
                        for c in state.table_state.columns
                        if c.title not in p.target_columns
                    ]
                case RemovePapersParams():
                    # capture cells that will be removed
                    for col in state.table_state.columns:
                        for cell in col.values:
                            if cell.paper_id in p.paper_ids:
                                orphaned_cells.append(cell)
                    # Remove paper_ids
                    state.table_state.paper_ids = [
                        pid
                        for pid in state.table_state.paper_ids
                        if pid not in p.paper_ids
                    ]
                    # Remove corresponding cell entries
                    for col in state.table_state.columns:
                        col.values = [
                            cell
                            for cell in col.values
                            if cell.paper_id not in p.paper_ids
                        ]
                case UndoParams():
                    if not state.history:
                        raise ToolError("No history to undo.")
                    last_json = state.history.pop()
                    state.table_state = TableData.model_validate_json(last_json)
                case GetTableParams():
                    pass
                case _:
                    raise ToolError(f"Unknown action {p.action!r}")
        except Exception:
            if mutating:
                state.history = prior_history
                state.table_state = TableData.model_validate_json(prior_table_json)
            raise

        # compute missing or empty cell values per column
        state.table_state.missing_cells = {}
        for col in state.table_state.columns:
            missing = []
            pid_to_cell = {c.paper_id: c for c in col.values}
            for pid in state.table_state.paper_ids:
                if pid not in pid_to_cell or pid_to_cell[pid].value == "":
                    missing.append(pid)
            state.table_state.missing_cells[col.title] = missing

        # return via payload, keep orphaned_cells in payload
        payload = TableResponsePayload(
            table=state.table_state,
            orphaned_cells=orphaned_cells,
        )
        return payload.model_dump_json()

    tool_def = ToolDef(execute, name="table_editor")
    if not enable_submission:
        # Remove the 'submit' part of the description
        lines = tool_def.parameters.properties["action"].description.split("\n")
        tool_def.parameters.properties["action"].description = "\n".join(
            line for line in lines if not line.startswith('- "submit"')
        )
        tool_def.parameters.properties["action"].enum.remove("submit")

    return tool_def.as_tool()