patch_document_upload.py

"""Patch: extend dataset uploads to support office/binary documents
(PDF, images, Excel, Word, PPT, OpenDocument, text) and extract their text
so the agent can understand the file content even on text-only LLM backends.
"""

import ast
import os

FILE = "/app/backend/dataset_uploads.py"

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

orig = c

# ---- 1. Expand allowed extensions ----
old_ext = 'ALLOWED_DATASET_EXTENSIONS = {"csv", "json", "jsonl"}'
new_ext = (
    'ALLOWED_DATASET_EXTENSIONS = {\n'
    '    "csv", "json", "jsonl",\n'
    '    "pdf",\n'
    '    "png", "jpg", "jpeg", "webp", "bmp", "tiff", "tif", "gif",\n'
    '    "xlsx", "xlsm", "xls",\n'
    '    "docx", "doc",\n'
    '    "pptx",\n'
    '    "odt", "ods", "odp",\n'
    '    "txt", "md", "rtf", "log",\n'
    '}\n'
    'BINARY_EXTRACTION_EXTS = {\n'
    '    "pdf",\n'
    '    "png", "jpg", "jpeg", "webp", "bmp", "tiff", "tif", "gif",\n'
    '    "xlsx", "xlsm", "xls",\n'
    '    "docx", "doc",\n'
    '    "pptx",\n'
    '    "odt", "ods", "odp",\n'
    '    "rtf",\n'
    '}\n'
)
if old_ext in c:
    c = c.replace(old_ext, new_ext)
    print("OK: expanded ALLOWED_DATASET_EXTENSIONS")

else:
    print("WARN: extension line not found")

# ---- 1b. Update the legacy error message so no stale message remains ----
old_msg = 'detail="Only .csv, .json, and .jsonl dataset files are supported."'
new_msg = 'detail="Không hỗ trợ định dạng này. Các định dạng hỗ trợ: CSV, JSON, JSONL, PDF, ảnh, Excel (XLSX/XLSM/XLS), Word (DOCX/DOC), PowerPoint (PPTX), OpenDocument (ODT/ODS/ODP), và văn bản (TXT/MD/RTF/LOG)."'
if old_msg in c:
    c = c.replace(old_msg, new_msg)
    print("OK: updated legacy error message")
else:
    print("WARN: legacy dataset message not found (already updated?)")


# ---- 2. Import the document parser ----
if "document_parser" not in c:
    c = c.replace(
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub import HfApi\n"
        "from document_parser import (\n"
        "    extract_document_text,\n"
        "    format_uploaded_document_context,\n"
        ")\n",
    )
    print("OK: imported document_parser")
else:
    print("OK: document_parser already imported")

# ---- 3. Extract text inside push_dataset_upload_to_hub ----
# Find the return statement that constructs the DatasetUpload and inject
# extraction before it.
marker = "    return DatasetUpload(\n"
if marker in c and "extract_document_text" in c and "extracted_text" not in c:
    extraction_block = (
        "    # Extract plain-text content so the agent can understand the file\n"
        "    # even when the LLM backend has no native attachment support.\n"
        "    extracted_text, extraction_warning = \"\", \"\"\n"
        "    if dataset_format in BINARY_EXTRACTION_EXTS:\n"
        "        try:\n"
        "            await asyncio.to_thread(upload.file.seek, 0)\n"
        "            raw_bytes = await asyncio.to_thread(upload.file.read)\n"
        "            extracted_text, extraction_warning = await asyncio.to_thread(\n"
        "                extract_document_text, safe_filename, raw_bytes\n"
        "            )\n"
        "            await asyncio.to_thread(upload.file.seek, 0)\n"
        "        except Exception as exc:  # never let extraction break the upload\n"
        "            extraction_warning = f\"Could not extract text: {exc}\"\n"
        "            extracted_text = \"\"\n"
        "        await asyncio.to_thread(upload.file.seek, 0)\n"
        "\n"
    )
    c = c.replace(marker, extraction_block + marker)
    print("OK: inserted extraction block before DatasetUpload return")
else:
    print("WARN: could not insert extraction block (marker or state issue)")

# ---- 4. Thread the extracted fields into the DatasetUpload dataclass ----
# Add fields to the dataclass
if "extracted_text: str" not in c:
    c = c.replace(
        "    load_dataset_snippet: str\n",
        "    load_dataset_snippet: str\n"
        "    extracted_text: str = \"\"\n"
        "    extraction_warning: str = \"\"\n",
    )
    print("OK: added dataclass fields")
else:
    print("OK: dataclass fields present")

# Add to response_payload
if '"extracted_text": self.extracted_text' not in c:
    c = c.replace(
        '            "load_dataset_snippet": self.load_dataset_snippet,\n',
        '            "load_dataset_snippet": self.load_dataset_snippet,\n'
        '            "extracted_text": self.extracted_text,\n'
        '            "extraction_warning": self.extraction_warning,\n',
    )
    print("OK: response_payload includes extracted_text")
else:
    print("OK: response_payload present")

# Add to the final constructor call
if "extracted_text=extracted_text" not in c:
    c = c.replace(
        "        load_dataset_snippet=snippet,\n"
        "    )",
        "        load_dataset_snippet=snippet,\n"
        "        extracted_text=extracted_text,\n"
        "        extraction_warning=extraction_warning,\n"
        "    )",
    )
    print("OK: constructor passes extracted fields")
else:
    print("OK: constructor present")

try:
    ast.parse(c)
    with open(FILE, "w") as f:
        f.write(c)
    print("OK: dataset_uploads.py written & syntax valid")
except SyntaxError as e:
    print(f"FAIL: dataset_uploads.py syntax error: {e}")

# ---- Also patch models.py DatasetUploadResponse.format literal ----
MODELS_FILE = "/app/backend/models.py"
with open(MODELS_FILE) as f:
    mc = f.read()

if 'Literal["csv", "json", "jsonl"]' in mc:
    mc = mc.replace(
        'format: Literal["csv", "json", "jsonl"]',
        'format: str',
    )
    # add new fields
    mc = mc.replace(
        "    load_dataset_snippet: str\n",
        "    load_dataset_snippet: str\n"
        "    extracted_text: str = \"\"\n"
        "    extraction_warning: str = \"\"\n",
    )
    with open(MODELS_FILE, "w") as f:
        f.write(mc)
    print("OK: models.py DatasetUploadResponse widened")
else:
    print("OK: models.py already patched")

# ---- Patch routes/agent.py to use the richer doc context note ----
AGENT_FILE = "/app/backend/routes/agent.py"
with open(AGENT_FILE) as f:
    ac = f.read()

if "format_uploaded_document_context" not in ac:
    # Update import
    old_imp = "    dataset_context_note,\n    push_dataset_upload_to_hub,\n"
    new_imp = (
        "    format_uploaded_document_context,\n"
        "    push_dataset_upload_to_hub,\n"
    )
    if old_imp in ac:
        ac = ac.replace(old_imp, new_imp)
        print("OK: agent.py import updated")
    else:
        # Try removing the dataset_context_note import line entirely
        ac = ac.replace("    dataset_context_note,\n", "")
        ac = ac.replace(
            "from dataset_uploads import (\n    push_dataset_upload_to_hub,\n)",
            "from dataset_uploads import (\n"
            "    format_uploaded_document_context,\n"
            "    push_dataset_upload_to_hub,\n)",
        )
        print("OK: agent.py import swapped (fallback)")

    # Update usage
    old_use = 'Message(role="user", content=dataset_context_note(uploaded))'
    new_use = (
        'Message(role="user", content=format_uploaded_document_context(\n'
        "            filename=uploaded.original_filename,\n"
        "            stored_filename=uploaded.filename,\n"
        "            repo_id=uploaded.repo_id,\n"
        "            path_in_repo=uploaded.path_in_repo,\n"
        "            hub_url=uploaded.hub_url,\n"
        "            size_bytes=uploaded.size_bytes,\n"
        "            file_format=uploaded.format,\n"
        "            extracted_text=uploaded.extracted_text,\n"
        "            warning=uploaded.extraction_warning,\n"
        "        ))"
    )
    if old_use in ac:
        ac = ac.replace(old_use, new_use)
        print("OK: agent.py usage updated")
    else:
        print("WARN: agent.py usage line not found")

    try:
        ast.parse(ac)
        with open(AGENT_FILE, "w") as f:
            f.write(ac)
        print("OK: agent.py written & syntax valid")
    except SyntaxError as e:
        print(f"FAIL: agent.py syntax error: {e}")
else:
    print("OK: agent.py already patched")

print("DONE: document upload backend patch complete")