brain/devtools/attachments.py

"""Attached files for the developer assistant. Master prompt 6.

Limits from the spec: text or code only, up to 1 MB each, at most 20 files. On top of that, a
total budget for what is injected into one prompt — twenty 1 MB files would blow through the
free tier's per-minute token allowance in a single turn, so past the budget each file is cut and
the model is told it was.
"""

import re
from dataclasses import dataclass

MAX_FILE_BYTES = 1_000_000
MAX_FILES = 20
#: Characters of attached text injected into one turn. Roughly 50k tokens.
PROMPT_BUDGET_CHARS = 200_000

#: Extensions whose content is worth sending. Anything else is refused at the door rather than
#: sent as mojibake.
TEXT_EXTENSIONS = frozenset(
    ".php .js .jsx .ts .tsx .mjs .cjs .py .rb .go .rs .java .kt .kts .swift .c .h .cpp .hpp .cs "
    ".html .htm .css .scss .sass .less .vue .svelte .json .yaml .yml .toml .ini .env.example .xml "
    ".sql .md .txt .sh .bash .ps1 .bat .dockerfile .gradle .properties .twig .blade.php .lock "
    ".graphql .prisma .csv .conf .htaccess .nginx".split()
)

_LANGUAGE = {
    ".php": "php",
    ".js": "javascript",
    ".jsx": "jsx",
    ".ts": "typescript",
    ".tsx": "tsx",
    ".py": "python",
    ".json": "json",
    ".yaml": "yaml",
    ".yml": "yaml",
    ".sql": "sql",
    ".css": "css",
    ".scss": "scss",
    ".html": "html",
    ".vue": "vue",
    ".sh": "bash",
    ".java": "java",
    ".kt": "kotlin",
    ".go": "go",
    ".rs": "rust",
    ".rb": "ruby",
    ".md": "markdown",
}

_FENCE = re.compile(r"```[\w+-]*\n(.*?)```", re.DOTALL)


class AttachmentError(ValueError):
    pass


@dataclass(frozen=True)
class Attachment:
    name: str
    content: str


def is_text_file(name: str) -> bool:
    lower = name.lower()
    base = lower.rsplit("/", 1)[-1]
    if base in {
        "dockerfile",
        "makefile",
        ".htaccess",
        "composer.json",
        "package.json",
        ".gitignore",
    }:
        return True
    return any(lower.endswith(ext) for ext in TEXT_EXTENSIONS)


def validate(files: list[Attachment]) -> list[Attachment]:
    if len(files) > MAX_FILES:
        raise AttachmentError(f"at most {MAX_FILES} files per message")
    for item in files:
        if not is_text_file(item.name):
            raise AttachmentError(f"{item.name}: only text and code files can be attached")
        if len(item.content.encode("utf-8")) > MAX_FILE_BYTES:
            raise AttachmentError(f"{item.name}: larger than 1 MB")
        if "\x00" in item.content:
            raise AttachmentError(f"{item.name}: looks like a binary file")
    return files


def language_for(name: str) -> str:
    lower = name.lower()
    if lower.endswith(".blade.php"):
        return "blade"
    for ext, language in _LANGUAGE.items():
        if lower.endswith(ext):
            return language
    return ""


def render_for_prompt(files: list[Attachment], tree: list[str] | None = None) -> str:
    """The attachments as one block of the user's message, inside the total budget."""
    if not files and not tree:
        return ""

    parts: list[str] = []
    if tree:
        shown = tree[:400]
        more = f"\n… and {len(tree) - len(shown)} more" if len(tree) > len(shown) else ""
        parts.append("PROJECT TREE\n" + "\n".join(shown) + more)

    remaining = PROMPT_BUDGET_CHARS
    per_file = PROMPT_BUDGET_CHARS // max(len(files), 1)
    for item in files:
        budget = min(per_file, remaining)
        body = item.content
        note = ""
        if len(body) > budget:
            body = body[:budget]
            note = f"\n[… truncated: {len(item.content) - budget} more characters not shown]"
        remaining -= len(body)
        fence = language_for(item.name)
        parts.append(f"FILE: {item.name}\n```{fence}\n{body}\n```{note}")

    return "ATTACHED FOR THIS MESSAGE\n\n" + "\n\n".join(parts)


def pasted_code(message: str) -> str:
    """Code he pasted into the message itself: the fenced blocks, or failing that, all of it."""
    blocks = _FENCE.findall(message)
    return "\n\n".join(block.strip() for block in blocks) if blocks else message.strip()


def resolve(target: str | None, attachments: dict[str, str], message: str) -> tuple[str, str]:
    """Which code a dev tool should look at: (label, content).

    A target naming an attachment wins — exactly, then by file name. Otherwise it is the code
    pasted into the message, and failing that the only attachment if there is exactly one.
    """
    if target:
        wanted = target.strip()
        if wanted in attachments:
            return wanted, attachments[wanted]
        by_name = {name.rsplit("/", 1)[-1]: name for name in attachments}
        if wanted in by_name:
            return by_name[wanted], attachments[by_name[wanted]]

    code = pasted_code(message)
    if _FENCE.search(message) or not attachments:
        return "pasted code", code
    if len(attachments) == 1:
        name, content = next(iter(attachments.items()))
        return name, content
    return "pasted code", code