"""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")