""" API routes for the ingestion pipeline """ import os import asyncio import uuid import time import json from typing import List, Dict, Any, Optional from fastapi import APIRouter, Form, File, UploadFile, HTTPException, BackgroundTasks, Request from api.models import UploadResponse, JobStatusResponse, HealthResponse, FilesListResponse, ChunksResponse from services.ingestion_service import IngestionService from services.maverick_captioner import _normalize_caption from utils.logger import get_logger logger = get_logger("INGESTION_ROUTES", __name__) # Create router router = APIRouter() # Global services (will be injected) rag = None embedder = None captioner = None ingestion_service = None def initialize_services(rag_store, embedder_client, captioner_client): """Initialize services""" global rag, embedder, captioner, ingestion_service rag = rag_store embedder = embedder_client captioner = captioner_client ingestion_service = IngestionService(rag_store, embedder_client, captioner_client) @router.get("/health", response_model=HealthResponse) async def health(): """Health check endpoint""" mongodb_connected = rag is not None return HealthResponse( ok=mongodb_connected, mongodb_connected=mongodb_connected ) @router.post("/upload", response_model=UploadResponse) async def upload_files( request: Request, background_tasks: BackgroundTasks, user_id: str = Form(...), project_id: str = Form(...), files: List[UploadFile] = File(...), replace_filenames: Optional[str] = Form(None), rename_map: Optional[str] = Form(None), ): """ Upload and process files This endpoint mirrors the main system's upload functionality exactly. """ if not rag: raise HTTPException(500, detail="MongoDB connection not available") job_id = str(uuid.uuid4()) # File limits (same as main system) max_files = int(os.getenv("MAX_FILES_PER_UPLOAD", "15")) max_mb = int(os.getenv("MAX_FILE_MB", "50")) if len(files) > max_files: raise HTTPException(400, detail=f"Too many files. Max {max_files} allowed per upload.") # Parse replace/rename directives (same as main system) replace_set = set() try: if replace_filenames: replace_set = set(json.loads(replace_filenames)) except Exception: pass rename_dict: Dict[str, str] = {} try: if rename_map: rename_dict = json.loads(rename_map) except Exception: pass # Preload files (same as main system) preloaded_files = [] for uf in files: raw = await uf.read() if len(raw) > max_mb * 1024 * 1024: raise HTTPException(400, detail=f"{uf.filename} exceeds {max_mb} MB limit") eff_name = rename_dict.get(uf.filename, uf.filename) preloaded_files.append((eff_name, raw)) # Initialize job status (same as main system) from app import app app.state.jobs[job_id] = { "created_at": time.time(), "total": len(preloaded_files), "completed": 0, "status": "processing", "last_error": None, } # Background processing (mirrors main system exactly) async def _process_all(): for idx, (fname, raw) in enumerate(preloaded_files, start=1): try: # Handle file replacement (same as main system) if fname in replace_set: try: rag.db["chunks"].delete_many({"user_id": user_id, "project_id": project_id, "filename": fname}) rag.db["files"].delete_many({"user_id": user_id, "project_id": project_id, "filename": fname}) logger.info(f"[{job_id}] Replaced prior data for {fname}") except Exception as de: logger.warning(f"[{job_id}] Replace delete failed for {fname}: {de}") logger.info(f"[{job_id}] ({idx}/{len(preloaded_files)}) Parsing {fname} ({len(raw)} bytes)") # Extract pages (same as main system) from helpers.pages import _extract_pages pages = _extract_pages(fname, raw) # Process images with captions (same as main system) num_imgs = sum(len(p.get("images", [])) for p in pages) captions = [] if num_imgs > 0: for p in pages: caps = [] for im in p.get("images", []): try: cap = captioner.caption_image(im) caps.append(cap) except Exception as e: logger.warning(f"[{job_id}] Caption error in {fname}: {e}") captions.append(caps) else: captions = [[] for _ in pages] # Merge captions into text (same as main system) for p, caps in zip(pages, captions): if caps: normalized = [ _normalize_caption(c) for c in caps if c ] if normalized: p["text"] = (p.get("text", "") + "\n\n" + "\n".join([f"[Image] {c}" for c in normalized])).strip() # Build cards (same as main system) from utils.ingestion.chunker import build_cards_from_pages cards = await build_cards_from_pages(pages, filename=fname, user_id=user_id, project_id=project_id) logger.info(f"[{job_id}] Built {len(cards)} cards for {fname}") # Generate embeddings (same as main system) embeddings = embedder.embed([c["content"] for c in cards]) for c, vec in zip(cards, embeddings): c["embedding"] = vec # Store in MongoDB (same as main system) rag.store_cards(cards) # Create file summary (same as main system) from utils.service.summarizer import cheap_summarize full_text = "\n\n".join(p.get("text", "") for p in pages) file_summary = await cheap_summarize(full_text, max_sentences=6) rag.upsert_file_summary(user_id=user_id, project_id=project_id, filename=fname, summary=file_summary) logger.info(f"[{job_id}] Completed {fname}") # Update job progress (same as main system) job = app.state.jobs.get(job_id) if job: job["completed"] = idx job["status"] = "processing" if idx < job.get("total", 0) else "completed" except Exception as e: logger.error(f"[{job_id}] Failed processing {fname}: {e}") job = app.state.jobs.get(job_id) if job: job["last_error"] = str(e) job["completed"] = idx finally: await asyncio.sleep(0) # Finalize job (same as main system) logger.info(f"[{job_id}] Ingestion complete for {len(preloaded_files)} files") job = app.state.jobs.get(job_id) if job: job["status"] = "completed" background_tasks.add_task(_process_all) return UploadResponse( job_id=job_id, status="processing", total_files=len(preloaded_files) ) @router.get("/upload/status", response_model=JobStatusResponse) async def upload_status(job_id: str): """Get upload job status""" from app import app job = app.state.jobs.get(job_id) if not job: raise HTTPException(404, detail="Job not found") progress = (job["completed"] / job["total"]) * 100 if job["total"] > 0 else 0 return JobStatusResponse( job_id=job_id, status=job["status"], total=job["total"], completed=job["completed"], progress=progress, last_error=job.get("last_error"), created_at=job["created_at"] ) @router.get("/files", response_model=FilesListResponse) async def list_files(user_id: str, project_id: str): """List files for a project (compatible with main system)""" if not rag: raise HTTPException(500, detail="MongoDB connection not available") files = rag.list_files(user_id, project_id) return FilesListResponse( files=[{"filename": f["filename"], "summary": f["summary"]} for f in files], filenames=[f["filename"] for f in files] ) @router.get("/files/chunks", response_model=ChunksResponse) async def get_file_chunks(user_id: str, project_id: str, filename: str, limit: int = 20): """Get chunks for a specific file (compatible with main system)""" if not rag: raise HTTPException(500, detail="MongoDB connection not available") chunks = rag.get_file_chunks(user_id, project_id, filename, limit) return ChunksResponse(chunks=chunks)