Spaces:
Sleeping
Sleeping
File size: 9,593 Bytes
f4cfd3f b76768d f4cfd3f b76768d f4cfd3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 |
import os
import uuid
import json
import shutil
from pathlib import Path
from typing import List, Dict, Any, Optional
from fastapi import UploadFile, HTTPException
import time
# Document handling libraries
import pypdf
import docx2txt
# Create data directories if they don't exist
DATA_DIR = Path("./data")
DOCUMENTS_DIR = DATA_DIR / "documents"
DOCUMENTS_DIR.mkdir(parents=True, exist_ok=True)
METADATA_FILE = DATA_DIR / "document_metadata.json"
# Hard cap on upload size — an unbounded file would fill the container's
# ephemeral disk / RAM and take the service down. 10 MB covers real docs.
MAX_UPLOAD_BYTES = 10 * 1024 * 1024
# Derive the saved-file extension from the validated content type, never from
# the client-supplied filename (which is attacker-controlled).
EXTENSION_FOR_CONTENT_TYPE = {
"application/pdf": ".pdf",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
"text/plain": ".txt",
"text/markdown": ".md",
}
class DocumentProcessor:
"""
Handles document processing including:
- Uploading and storing documents
- Extracting text from different document formats
- Maintaining document metadata
"""
def __init__(self):
self.supported_formats = {
"application/pdf": self._extract_pdf_text,
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": self._extract_docx_text,
"text/plain": self._extract_text_file,
"text/markdown": self._extract_text_file,
}
self._initialize_metadata()
def _initialize_metadata(self):
"""Initialize metadata storage"""
if not METADATA_FILE.exists():
with open(METADATA_FILE, "w") as f:
json.dump({}, f)
async def process_document(self, file: UploadFile, title: Optional[str] = None, source: str = "user-upload") -> str:
"""
Process an uploaded document:
1. Extract text based on file type
2. Save document metadata
3. Return document ID for further processing
"""
# Generate a unique ID for the document
doc_id = str(uuid.uuid4())
# Determine document format and extract text
content_type = file.content_type
if content_type not in self.supported_formats:
raise HTTPException(
status_code=400,
detail=f"Unsupported file format: {content_type}. Supported formats: {', '.join(self.supported_formats.keys())}"
)
# Save the original file. Extension comes from the validated content
# type, NOT the client filename. Stream in chunks and abort if the
# upload exceeds MAX_UPLOAD_BYTES so a huge file can't exhaust disk/RAM.
extension = EXTENSION_FOR_CONTENT_TYPE.get(content_type, "")
file_path = DOCUMENTS_DIR / f"{doc_id}{extension}"
bytes_written = 0
try:
with open(file_path, "wb") as f:
while True:
chunk = await file.read(1024 * 1024) # 1 MB at a time
if not chunk:
break
bytes_written += len(chunk)
if bytes_written > MAX_UPLOAD_BYTES:
f.close()
file_path.unlink(missing_ok=True)
raise HTTPException(
status_code=413,
detail="File too large — uploads are limited to 10 MB.",
)
f.write(chunk)
except HTTPException:
raise
except Exception:
file_path.unlink(missing_ok=True)
raise
# Reset file position so the extractor below can re-read from the start
await file.seek(0)
# Extract text
text = await self.supported_formats[content_type](file)
# If title is not provided, use the filename without extension
if not title:
title = Path(file.filename).stem
# Save text to a file
text_file_path = DOCUMENTS_DIR / f"{doc_id}.txt"
with open(text_file_path, "w", encoding="utf-8") as f:
f.write(text)
# Update metadata
metadata = {
"id": doc_id,
"title": title,
"filename": file.filename,
"content_type": content_type,
"source": source,
"length": len(text),
"upload_time": time.time(),
"original_file": str(file_path),
"text_file": str(text_file_path)
}
self._save_metadata(doc_id, metadata)
return doc_id
async def _extract_pdf_text(self, file: UploadFile) -> str:
"""Extract text from PDF files"""
temp_path = DOCUMENTS_DIR / f"temp_{uuid.uuid4()}.pdf"
try:
# Save to temporary file
with open(temp_path, "wb") as f:
shutil.copyfileobj(file.file, f)
# Process PDF using pypdf
text = ""
with open(temp_path, "rb") as f:
pdf = pypdf.PdfReader(f)
for page in pdf.pages:
text += page.extract_text() + "\n\n"
return text
finally:
# Clean up temp file
if temp_path.exists():
os.remove(temp_path)
async def _extract_docx_text(self, file: UploadFile) -> str:
"""Extract text from DOCX files"""
temp_path = DOCUMENTS_DIR / f"temp_{uuid.uuid4()}.docx"
try:
# Save to temporary file
with open(temp_path, "wb") as f:
shutil.copyfileobj(file.file, f)
# Process DOCX using docx2txt
text = docx2txt.process(temp_path)
return text
finally:
# Clean up temp file
if temp_path.exists():
os.remove(temp_path)
async def _extract_text_file(self, file: UploadFile) -> str:
"""Extract text from plain text files"""
content = await file.read()
try:
return content.decode("utf-8")
except UnicodeDecodeError:
# Try different encoding if UTF-8 fails
return content.decode("latin-1")
def _save_metadata(self, doc_id: str, metadata: Dict[str, Any]):
"""Save document metadata to the metadata file"""
try:
# Load existing metadata
with open(METADATA_FILE, "r") as f:
all_metadata = json.load(f)
# Add new metadata
all_metadata[doc_id] = metadata
# Save updated metadata
with open(METADATA_FILE, "w") as f:
json.dump(all_metadata, f, indent=2)
except Exception as e:
print(f"Error saving metadata: {e}")
raise HTTPException(status_code=500, detail=f"Error saving document metadata: {str(e)}")
def list_documents(self) -> List[Dict[str, Any]]:
"""Get a list of all document metadata"""
try:
with open(METADATA_FILE, "r") as f:
all_metadata = json.load(f)
# Return a list of metadata objects
return list(all_metadata.values())
except Exception as e:
print(f"Error listing documents: {e}")
return []
def delete_document(self, doc_id: str) -> bool:
"""Delete a document and its metadata"""
try:
# Load existing metadata
with open(METADATA_FILE, "r") as f:
all_metadata = json.load(f)
# Check if document exists
if doc_id not in all_metadata:
return False
# Get metadata for the document
metadata = all_metadata[doc_id]
# Delete the original file if it exists
original_file = Path(metadata.get("original_file", ""))
if original_file.exists():
os.remove(original_file)
# Delete the text file if it exists
text_file = Path(metadata.get("text_file", ""))
if text_file.exists():
os.remove(text_file)
# Remove from metadata
del all_metadata[doc_id]
# Save updated metadata
with open(METADATA_FILE, "w") as f:
json.dump(all_metadata, f, indent=2)
return True
except Exception as e:
print(f"Error deleting document: {e}")
return False
def get_document_text(self, doc_id: str) -> Optional[str]:
"""Get the extracted text for a document"""
try:
# Load metadata
with open(METADATA_FILE, "r") as f:
all_metadata = json.load(f)
# Check if document exists
if doc_id not in all_metadata:
return None
# Get text file path
text_file = Path(all_metadata[doc_id].get("text_file", ""))
if not text_file.exists():
return None
# Read text file
with open(text_file, "r", encoding="utf-8") as f:
return f.read()
except Exception as e:
print(f"Error getting document text: {e}")
return None
|