feat(documents): async ingestion and state machine (#10)

* feat(documents): add upload endpoint and async state machine

* test(documents): cover resumability and per-stage idempotency

app/main.py

@@ -8,6 +8,8 @@ from app.routers import auth,hubs
 from app.deps import get_current_user
 from app.models import User
 from app.services.gemini import AiUnavailable
+from app.routers import auth, hubs, documents
+
 
 
 logger = logging.getLogger("uvicorn.error")
@@ -22,6 +24,8 @@ def create_app() -> FastAPI:
 
     app.include_router(auth.router, prefix="/api/v1")
     app.include_router(hubs.router, prefix="/api/v1")
+    app.include_router(documents.router, prefix="/api/v1")
+
 
     @app.middleware("http")
     async def catch_unhandled(request: Request, call_next):

app/routers/documents.py

@@ -0,0 +1,160 @@
+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)
+):
+    # 1. Verify hub ownership
+    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.")
+
+    # 2. Validate MIME type and size
+    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.")
+    
+    # Read bytes (limit to 10MB)
+    contents = await file.read()
+    if len(contents) > 10 * 1024 * 1024:
+        raise AppError(413, "DOCUMENT_TOO_LARGE", "File exceeds 10MB.")
+
+    # 3. Calculate checksum
+    checksum = hashlib.sha256(contents).hexdigest()
+
+    # 4. Idempotency hash (hub_id + filename + size + checksum)
+    hash_input = f"{hub_id}{file.filename}{len(contents)}{checksum}"
+    req_hash = hashlib.sha256(hash_input.encode()).hexdigest()
+
+    # 5. Claim key
+    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)
+
+    # 6. Check for duplicate upload
+    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.")
+
+    # 7. Save document
+    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)
+
+    # 8. Build response and complete key
+    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)
+):
+    # 1. Fetch document with tenant isolation
+    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.")
+
+    # 2. Idempotency check: If the same key is used, return the current state without advancing
+    if key and doc.process_claim and str(doc.process_claim) == str(key):
+        return DocumentResponse.model_validate(doc)
+
+    # 3. Claim the stage
+    if key:
+        doc.process_claim = key
+        db.commit() # Commit the claim immediately so concurrent requests see it
+
+    # 4. Advance state machine
+    if doc.status == "pending":
+        doc.status = "chunking"
+        db.commit()
+        
+    elif doc.status == "chunking":
+        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  # Placeholder
+                ))
+                
+            doc.file_bytes = None  # Delete raw bytes!
+            doc.status = "embedding"
+            db.commit()
+        except Exception as e:
+            doc.status = "failed"
+            doc.error_message = "PDF parsing failed: " + str(e)
+            db.commit()
+            
+    elif doc.status == "embedding":
+        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
+                
+                doc.status = "ready"
+                doc.process_claim = None # Clear claim on success
+                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)

app/schemas.py

@@ -77,4 +77,16 @@ class HubCreate(BaseModel):
     content: str = Field(default="", max_length=20000)      
     documents: list[DocumentSummary] = []    
 
-    
+class DocumentResponse(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    id: UUID
+    hub_id: UUID
+    origin: str
+    filename: str
+    mime_type: str
+    size_bytes: int
+    status: str
+    error_message: str | None
+    chunk_count: int = 0
+    created_at: datetime
+    updated_at: datetime    

tests/test_ingestion.py

@@ -0,0 +1,71 @@
+import pytest
+from uuid import uuid4
+from unittest.mock import patch
+from app.models import Document, Chunk
+
+def test_upload_and_process_pdf(client, db_session):
+    # Setup user and hub
+    client.post("/api/v1/auth/signup", json={"email": "doc@test.com", "password": "Password123"})
+    token = client.post("/api/v1/auth/login", json={"email": "doc@test.com", "password": "Password123"}).json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    
+    # Create a hub
+    key = str(uuid4())
+    hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": key}, json={"name": "PDF Hub"})
+    hub_id = hub_res.json()["id"]
+
+    # Upload a fake text file
+    file_content = b"This is a test document. It has multiple sentences to chunk."
+    upload_key = str(uuid4())
+    files = {"file": ("test.txt", file_content, "text/plain")}
+    
+    upload_res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": upload_key}, files=files)
+    assert upload_res.status_code == 202
+    doc_id = upload_res.json()["id"]
+    assert upload_res.json()["status"] == "pending"
+
+    # Process Stage 1 (chunking)
+    with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
+        res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
+        assert res1.json()["status"] == "chunking"
+
+        # Process Stage 2 (embedding)
+        res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
+        assert res2.json()["status"] == "embedding"
+        
+        # Process Stage 3 (ready)
+        res3 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
+        assert res3.json()["status"] == "ready"
+
+    # Verify chunks exist in DB
+    assert db_session.query(Chunk).count() > 0
+
+def test_process_is_idempotent(client, db_session):
+    # Setup user and hub
+    client.post("/api/v1/auth/signup", json={"email": "idem_doc@test.com", "password": "Password123"})
+    token = client.post("/api/v1/auth/login", json={"email": "idem_doc@test.com", "password": "Password123"}).json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    
+    hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Idem Doc Hub"})
+    hub_id = hub_res.json()["id"]
+
+    files = {"file": ("test.txt", b"Test text for idempotency.", "text/plain")}
+    upload_res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": str(uuid4())}, files=files)
+    doc_id = upload_res.json()["id"]
+
+    key1 = str(uuid4())
+    
+    with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
+        # 1. Advance to chunking
+        res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
+        assert res1.json()["status"] == "chunking"
+
+        # 2. Retry the exact same request with key1
+        res1_retry = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
+        # It should return "chunking", NOT "embedding"!
+        assert res1_retry.json()["status"] == "chunking"
+
+        # 3. Advance to embedding with a NEW key
+        key2 = str(uuid4())
+        res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key2})
+        assert res2.json()["status"] == "embedding"