app.py

"""
app.py
Gradio entry point
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import gradio as gr

from config import (
    DEFAULT_MAX_FILE_SIZE_MB,
    DEFAULT_MAX_TOKENS,
    IGNORE_DIRS_DEFAULT,
)
from core import build_context


def run(
    zip_file,
    ignored_dirs,
    output_format,
    max_tokens,
    max_size,
    include_hash,
    progress=gr.Progress(),
):

    progress(0.05, desc="Extracting ZIP")

    result = build_context(
        zip_path=zip_file,
        ignored_dirs=ignored_dirs,
        allowed_extensions=None,
        max_size_mb=max_size,
        output_format=output_format,
        max_tokens=max_tokens,
        include_hash=include_hash,
    )

    progress(0.8, desc="Writing output")

    temp = Path(tempfile.mkdtemp())

    downloads = []

    for i, chunk in enumerate(result["chunks"], start=1):

        if len(result["chunks"]) == 1:

            filename = (
                f"context.{result['extension']}"
            )

        else:

            filename = (
                f"context_part{i}.{result['extension']}"
            )

        file = temp / filename

        file.write_text(
            chunk,
            encoding="utf-8",
        )

        downloads.append(str(file))

    progress(1.0, desc="Done")

    return (
        result["statistics"],
        result["tree"],
        result["preview"],
        downloads,
    )


with gr.Blocks(title="Project Context Builder") as demo:

    gr.Markdown("# Project Context Builder")

    zip_file = gr.File(
        label="Project ZIP",
        file_types=[".zip"],
        type="filepath",
    )

    ignored = gr.Textbox(
        label="Ignored folders",
        value=IGNORE_DIRS_DEFAULT,
    )

    output = gr.Radio(
        ["Markdown", "XML"],
        value="Markdown",
        label="Output",
    )

    tokens = gr.Number(
        value=DEFAULT_MAX_TOKENS,
        label="Max Tokens (0 = unlimited)",
    )

    size = gr.Number(
        value=DEFAULT_MAX_FILE_SIZE_MB,
        label="Max File Size (MB)",
    )

    sha = gr.Checkbox(
        value=True,
        label="Include SHA256",
    )

    btn = gr.Button("Build")

    stats = gr.Markdown()

    tree = gr.Textbox(
        label="Project Tree",
        lines=15,
    )

    preview = gr.Textbox(
        label="Preview",
        lines=20,
    )

    downloads = gr.Files(
        label="Download",
    )

    btn.click(
        fn=run,
        inputs=[
            zip_file,
            ignored,
            output,
            tokens,
            size,
            sha,
        ],
        outputs=[
            stats,
            tree,
            preview,
            downloads,
        ],
    )

demo.queue()

if __name__ == "__main__":
    demo.launch()