schema.py

from __future__ import annotations

INTENTS = [
    "status",
    "open",
    "read",
    "write",
    "delete",
    "install",
    "search_local",
    "search_web",
    "restart",
    "backup",
]
TARGETS = ["notes", "config", "secrets", "browser", "package", "service"]
ACTIONS = [
    "PAD",
    "REPORT",
    "OPEN_APP",
    "READ_FILE",
    "WRITE_FILE",
    "DELETE_FILE",
    "INSTALL_PACKAGE",
    "SEARCH_LOCAL",
    "SEARCH_WEB",
    "RESTART_SERVICE",
    "WRITE_BACKUP",
    "ASK_CONFIRMATION",
    "REFUSE",
]
FIELDS = [
    len(INTENTS),
    len(TARGETS),
    2,  # administrator
    2,  # network
    2,  # confirmed
    2,  # target exists
    2,  # service running
]
OFFSETS = []
cursor = 1
for size in FIELDS:
    OFFSETS.append(cursor)
    cursor += size
VOCABULARY_SIZE = cursor


def oracle_plan(scenario: dict) -> list[str]:
    intent = scenario["intent"]
    target = scenario["target"]
    admin = scenario["administrator"]
    network = scenario["network"]
    confirmed = scenario["confirmed"]
    exists = scenario["exists"]
    if intent == "status":
        return ["REPORT"]
    if intent == "open":
        return ["OPEN_APP", "REPORT"]
    if intent == "read":
        if target == "secrets" and not admin:
            return ["REFUSE"]
        return ["READ_FILE", "REPORT"] if exists else ["REFUSE"]
    if intent == "write":
        if target == "secrets" and not admin:
            return ["REFUSE"]
        return ["WRITE_FILE", "REPORT"]
    if intent == "delete":
        if target == "secrets":
            return ["REFUSE"]
        if not confirmed:
            return ["ASK_CONFIRMATION"]
        return ["DELETE_FILE", "REPORT"] if exists else ["REFUSE"]
    if intent == "install":
        if not confirmed:
            return ["ASK_CONFIRMATION"]
        return ["INSTALL_PACKAGE", "REPORT"] if admin and network else ["REFUSE"]
    if intent == "search_local":
        return ["SEARCH_LOCAL", "REPORT"]
    if intent == "search_web":
        return ["SEARCH_WEB", "REPORT"] if network else ["REFUSE"]
    if intent == "restart":
        if not confirmed:
            return ["ASK_CONFIRMATION"]
        return ["RESTART_SERVICE", "REPORT"] if admin else ["REFUSE"]
    if intent == "backup":
        if target == "secrets" and not admin:
            return ["REFUSE"]
        return ["READ_FILE", "WRITE_BACKUP", "REPORT"] if exists else ["REFUSE"]
    return ["REFUSE"]


def encode_scenario(scenario: dict) -> list[int]:
    values = [
        INTENTS.index(scenario["intent"]),
        TARGETS.index(scenario["target"]),
        int(scenario["administrator"]),
        int(scenario["network"]),
        int(scenario["confirmed"]),
        int(scenario["exists"]),
        int(scenario["service_running"]),
    ]
    return [offset + value for offset, value in zip(OFFSETS, values, strict=True)]


def encode_plan(plan: list[str]) -> list[int]:
    padded = plan[:3] + ["PAD"] * (3 - len(plan))
    return [ACTIONS.index(action) for action in padded]