patch_multi_upload.py

"""Patch: allow multiple file uploads in one request.

Backend changes:
  1. Replace the single-file helper with a multi-file helper.
  2. Rewrite the upload route to accept up to 20 files and return a list of
     DatasetUploadResponse, uploading each file to the session's Hub dataset
     repo and injecting a context note for each.
"""

import ast
import re
import os

AGENT_FILE = "/app/backend/routes/agent.py"

with open(AGENT_FILE) as f:
    c = f.read()

# ---- 1. Replace single-file helper with multi-file helper ----
old_helper = '''def _dataset_upload_file_from_form(form: FormData) -> UploadFile:
    uploaded_files = [
        (key, value)
        for key, value in form.multi_items()
        if isinstance(value, UploadFile)
    ]
    if len(uploaded_files) != 1:
        raise HTTPException(
            status_code=400,
            detail="Upload exactly one dataset file.",
        )
    field_name, upload = uploaded_files[0]
    if field_name != "file":
        raise HTTPException(
            status_code=400,
            detail="Missing 'file' upload field.",
        )
    return upload
'''

new_helper = '''def _dataset_upload_files_from_form(form: FormData) -> list[UploadFile]:
    """Return all uploaded files from the multipart form (field 'file' or 'files')."""
    uploaded_files = [
        value
        for key, value in form.multi_items()
        if isinstance(value, UploadFile)
    ]
    if not uploaded_files:
        raise HTTPException(
            status_code=400,
            detail="No file uploads found in the request.",
        )
    return uploaded_files
'''

if old_helper in c:
    c = c.replace(old_helper, new_helper)
    print("OK: replaced helper with multi-file helper")
else:
    print("WARN: single-file helper not found (check indentation)")

# ---- 2. Rewrite the upload route to handle multiple files ----
# Replace the entire function based on its decorator + route path.
start_marker = '@router.post("/session/{session_id}/datasets"'
end_marker = '@router.patch("/session/{session_id}/yolo")'
start = c.find(start_marker)
end = c.find(end_marker)
if start == -1 or end == -1:
    print(f"FAIL: route markers not found (start={start}, end={end})")
    raise SystemExit(1)

new_route = '''@router.post("/session/{session_id}/datasets", response_model=list[DatasetUploadResponse])
async def upload_session_dataset(
    session_id: str,
    request: Request,
    user: dict = Depends(get_current_user),
) -> list[DatasetUploadResponse]:
    """Upload one or more document/dataset files to a private Hub dataset for this session."""
    files: list[UploadFile] = []
    try:
        _reject_oversize_dataset_upload(request)
        agent_session = await _check_session_access(session_id, user, request)
        if not agent_session or not agent_session.is_active:
            raise HTTPException(status_code=404, detail="Session not found")
        if agent_session.is_processing:
            raise HTTPException(
                status_code=409,
                detail="Cannot upload files while the agent is processing.",
            )
        if agent_session.session.pending_approval:
            raise HTTPException(
                status_code=409,
                detail="Resolve pending approvals before uploading files.",
            )

        hf_token = (
            resolve_hf_request_token(request, include_env_fallback=False)
            or _user_hf_token(user)
            or resolve_hf_request_token(request)
        )
        if not hf_token:
            raise HTTPException(
                status_code=401,
                detail="A Hugging Face token is required to upload files.",
            )

        form = await request.form(
            max_files=20,
            max_fields=20,
            max_part_size=MAX_DATASET_UPLOAD_BYTES,
        )
        files = _dataset_upload_files_from_form(form)
        if not files:
            raise HTTPException(status_code=400, detail="No files to upload.")

        hf_username = user.get("username") or agent_session.hf_username
        results: list[DatasetUploadResponse] = []
        for file in files:
            uploaded = await push_dataset_upload_to_hub(
                upload=file,
                session_id=session_id,
                hf_username=hf_username,
                hf_token=hf_token,
            )
            agent_session.session.context_manager.add_message(
                Message(role="user", content=format_uploaded_document_context(
                    filename=uploaded.original_filename,
                    stored_filename=uploaded.filename,
                    repo_id=uploaded.repo_id,
                    path_in_repo=uploaded.path_in_repo,
                    hub_url=uploaded.hub_url,
                    size_bytes=uploaded.size_bytes,
                    file_format=uploaded.format,
                    extracted_text=uploaded.extracted_text,
                    warning=uploaded.extraction_warning,
                ))
            )
            results.append(DatasetUploadResponse(**uploaded.response_payload()))
            logger.info(
                "Uploaded file %s to %s for session %s",
                uploaded.filename,
                uploaded.repo_id,
                session_id,
            )

        session_manager._touch(agent_session)
        await session_manager.persist_session_snapshot(agent_session)
        return results
    except HTTPException:
        raise
    except HfHubHTTPError as e:
        logger.warning(
            "Hub rejected file upload for session %s: status=%s request_id=%s",
            session_id,
            getattr(e.response, "status_code", None),
            getattr(e, "request_id", None),
        )
        raise _dataset_upload_hub_http_exception(e)
    except Exception:
        logger.exception("File upload failed for session %s", session_id)
        raise HTTPException(
            status_code=502,
            detail="File upload failed. Please try again.",
        )
    finally:
        for f in files:
            await f.close()


'''

c = c[:start] + new_route + c[end:]

with open(AGENT_FILE, "w") as f:
    f.write(c)

try:
    ast.parse(c)
    print("OK: agent.py written & syntax valid")
except SyntaxError as e:
    print(f"FAIL: agent.py syntax error: {e}")
    raise

print("DONE: multi-file upload backend patch complete")