Spaces:
Runtime error
Runtime error
File size: 6,326 Bytes
b13c030 cc78da3 b13c030 cc78da3 b13c030 cc78da3 b13c030 cc78da3 b13c030 cc78da3 b13c030 cc78da3 b13c030 cc78da3 b13c030 | 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 | import hashlib
import io
from uuid import UUID
from fastapi import APIRouter, Depends, File, UploadFile, Header, status
from sqlalchemy.orm import Session
from pypdf import PdfReader
from app.db import get_db
from app.deps import get_current_user
from app.models import User, Hub, Document, Chunk
from app.schemas import DocumentResponse
from app.errors import AppError
from app.idempotency import (
require_idempotency_key, purge_stale, claim_key,
complete_key, release_claim, replay_or_conflict
)
from app.services.embeddings import embed_texts
from app.chunking import chunk_text
router = APIRouter(prefix="/hubs/{hub_id}/documents", tags=["documents"])
@router.post("", response_model=DocumentResponse, status_code=status.HTTP_202_ACCEPTED)
async def upload_document(
hub_id: UUID,
file: UploadFile = File(...),
key: UUID = Depends(require_idempotency_key),
current_user: User = Depends(get_current_user),
db: Session = Depends(get_db)
):
hub = db.query(Hub).filter(Hub.id == hub_id, Hub.owner_id == current_user.id).first()
if not hub:
raise AppError(404, "HUB_NOT_FOUND", "Hub not found.")
allowed_mimes = ["application/pdf", "text/plain", "text/markdown"]
if file.content_type not in allowed_mimes:
raise AppError(415, "UNSUPPORTED_MIME_TYPE", "Only PDF, TXT, and MD are allowed.")
contents = await file.read()
if len(contents) > 10 * 1024 * 1024:
raise AppError(413, "DOCUMENT_TOO_LARGE", "File exceeds 10MB.")
checksum = hashlib.sha256(contents).hexdigest()
hash_input = f"{hub_id}{file.filename}{len(contents)}{checksum}"
req_hash = hashlib.sha256(hash_input.encode()).hexdigest()
purge_stale(db, current_user.id)
if not claim_key(db, current_user.id, key, req_hash, f"/api/v1/hubs/{hub_id}/documents"):
return replay_or_conflict(db, current_user.id, key, req_hash)
existing = db.query(Document).filter(Document.hub_id == hub_id, Document.checksum_sha256 == checksum).first()
if existing:
release_claim(db, current_user.id, key)
raise AppError(409, "DOCUMENT_ALREADY_EXISTS", "This file has already been uploaded to this hub.")
new_doc = Document(
hub_id=hub_id, owner_id=current_user.id, origin="upload",
filename=file.filename, mime_type=file.content_type, size_bytes=len(contents),
checksum_sha256=checksum, status="pending", file_bytes=contents
)
db.add(new_doc)
db.commit()
db.refresh(new_doc)
payload = DocumentResponse.model_validate(new_doc).model_dump(mode="json")
complete_key(db, current_user.id, key, 202, payload)
return payload
@router.post("/{doc_id}/process", response_model=DocumentResponse)
def process_document(
hub_id: UUID,
doc_id: UUID,
key: UUID | None = Header(default=None, alias="X-Idempotency-Key"),
current_user: User = Depends(get_current_user),
db: Session = Depends(get_db)
):
doc = db.query(Document).filter(
Document.id == doc_id, Document.hub_id == hub_id, Document.owner_id == current_user.id
).first()
if not doc:
raise AppError(404, "DOCUMENT_NOT_FOUND", "Document not found.")
if doc.status in ("ready", "failed"):
raise AppError(409, "DOCUMENT_NOT_PROCESSABLE", "Document is already complete or failed.")
if key and doc.process_claim and str(doc.process_claim) == str(key):
return DocumentResponse.model_validate(doc)
if key:
doc.process_claim = key
db.commit()
if doc.status == "pending":
rows = db.query(Document).filter(
Document.id == doc.id, Document.status == "pending"
).update({"status": "chunking"}, synchronize_session=False)
db.commit()
db.refresh(doc)
return DocumentResponse.model_validate(doc)
elif doc.status == "chunking":
rows = db.query(Document).filter(
Document.id == doc.id, Document.status == "chunking"
).update({"status": "embedding"}, synchronize_session=False)
db.commit()
if rows == 0:
db.refresh(doc)
return DocumentResponse.model_validate(doc)
try:
if doc.mime_type == "application/pdf":
reader = PdfReader(io.BytesIO(doc.file_bytes))
text = "\n".join(page.extract_text() for page in reader.pages if page.extract_text())
else:
text = doc.file_bytes.decode("utf-8", errors="ignore")
if not text.strip():
raise ValueError("No text extracted")
chunks_data = chunk_text(text)
for i, (chunk_text_str, tokens) in enumerate(chunks_data):
db.add(Chunk(
document_id=doc.id, hub_id=hub_id, owner_id=current_user.id,
chunk_index=i, text=chunk_text_str, token_count=tokens,
embedding=[0.0] * 768
))
doc.file_bytes = None
db.commit()
except Exception as e:
doc.status = "failed"
doc.error_message = "PDF parsing failed: " + str(e)
db.commit()
elif doc.status == "embedding":
rows = db.query(Document).filter(
Document.id == doc.id, Document.status == "embedding"
).update({"status": "ready"}, synchronize_session=False)
db.commit()
if rows == 0:
db.refresh(doc)
return DocumentResponse.model_validate(doc)
chunks = db.query(Chunk).filter(Chunk.document_id == doc.id).order_by(Chunk.chunk_index).all()
if not chunks:
doc.status = "failed"
doc.error_message = "No chunks found to embed."
db.commit()
else:
try:
texts = [c.text for c in chunks]
embeddings = embed_texts(texts)
for chunk, emb in zip(chunks, embeddings):
chunk.embedding = emb
db.commit()
except Exception as e:
doc.status = "failed"
doc.error_message = "Embedding failed: " + str(e)
db.commit()
db.refresh(doc)
return DocumentResponse.model_validate(doc) |