""" Docling Service - Document extraction and processing for the AI stack Integrates with: - AI Orchestrator (Ollama) for summarization - RunPod Whisper for audio transcription - Semantic Search for indexing extracted content """ import os import asyncio import tempfile import base64 import hashlib from pathlib import Path from datetime import datetime from typing import Optional, List, Dict, Any, Literal from enum import Enum import httpx from fastapi import FastAPI, HTTPException, UploadFile, File, Form, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, HTMLResponse from pydantic import BaseModel, HttpUrl # Docling imports from docling.document_converter import DocumentConverter from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import PdfPipelineOptions from docling.document_converter import PdfFormatOption # Config from environment AI_ORCHESTRATOR_URL = os.getenv("AI_ORCHESTRATOR_URL", "http://ai-orchestrator:8080") SEMANTIC_SEARCH_URL = os.getenv("SEMANTIC_SEARCH_URL", "http://semantic-search:8000") RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY", "") RUNPOD_WHISPER_ENDPOINT = os.getenv("RUNPOD_WHISPER_ENDPOINT", "lrtisuv8ixbtub") # Supported formats DOCUMENT_FORMATS = {".pdf", ".docx", ".pptx", ".xlsx", ".html", ".md", ".txt", ".epub"} IMAGE_FORMATS = {".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif"} AUDIO_FORMATS = {".mp3", ".wav", ".m4a", ".ogg", ".flac", ".webm"} app = FastAPI( title="Docling Service", description="Document extraction and processing service using Docling", version="1.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Initialize document converter with optimized settings pipeline_options = PdfPipelineOptions() pipeline_options.do_ocr = True pipeline_options.do_table_structure = True converter = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options) } ) # Track processing stats stats = { "documents_processed": 0, "pages_extracted": 0, "audio_transcribed": 0, "urls_fetched": 0, "errors": 0, } class OutputFormat(str, Enum): MARKDOWN = "markdown" JSON = "json" TEXT = "text" HTML = "html" class SummarizeStyle(str, Enum): CONCISE = "concise" DETAILED = "detailed" BULLET_POINTS = "bullet_points" TECHNICAL = "technical" ELI5 = "eli5" # Explain like I'm 5 class ExtractRequest(BaseModel): """Request to extract content from a URL or base64-encoded file""" url: Optional[HttpUrl] = None base64_content: Optional[str] = None filename: Optional[str] = None output_format: OutputFormat = OutputFormat.MARKDOWN summarize: bool = False summarize_style: SummarizeStyle = SummarizeStyle.CONCISE index_to_search: bool = False metadata: Optional[Dict[str, Any]] = None class TranscribeRequest(BaseModel): """Request to transcribe audio""" url: Optional[HttpUrl] = None base64_content: Optional[str] = None language: Optional[str] = None # Auto-detect if not specified summarize: bool = False summarize_style: SummarizeStyle = SummarizeStyle.CONCISE class BatchExtractRequest(BaseModel): """Batch extraction request""" items: List[ExtractRequest] class ExtractionResult(BaseModel): """Result of document extraction""" success: bool source: str content: Optional[str] = None format: OutputFormat metadata: Dict[str, Any] = {} summary: Optional[str] = None indexed: bool = False error: Optional[str] = None # ============== Helper Functions ============== def get_file_extension(filename: str) -> str: """Get lowercase file extension""" return Path(filename).suffix.lower() def generate_doc_id(source: str, content: str) -> str: """Generate a unique document ID""" hash_input = f"{source}:{content[:1000]}" return hashlib.sha256(hash_input.encode()).hexdigest()[:16] async def fetch_url_content(url: str) -> tuple[bytes, str]: """Fetch content from URL, return bytes and detected filename""" async with httpx.AsyncClient(follow_redirects=True, timeout=60) as client: resp = await client.get(url) resp.raise_for_status() # Try to get filename from headers or URL content_disposition = resp.headers.get("content-disposition", "") if "filename=" in content_disposition: filename = content_disposition.split("filename=")[1].strip('"\'') else: filename = url.split("/")[-1].split("?")[0] or "document" return resp.content, filename async def transcribe_audio_runpod(audio_data: bytes, language: Optional[str] = None) -> dict: """Transcribe audio using RunPod Whisper endpoint""" if not RUNPOD_API_KEY: raise HTTPException(status_code=500, detail="RunPod API key not configured") # Convert audio to base64 audio_base64 = base64.b64encode(audio_data).decode() payload = { "input": { "audio_base64": audio_base64, } } if language: payload["input"]["language"] = language async with httpx.AsyncClient(timeout=300) as client: # Submit job resp = await client.post( f"https://api.runpod.ai/v2/{RUNPOD_WHISPER_ENDPOINT}/run", headers={ "Authorization": f"Bearer {RUNPOD_API_KEY}", "Content-Type": "application/json", }, json=payload, ) result = resp.json() if "error" in result: raise HTTPException(status_code=500, detail=f"RunPod error: {result['error']}") job_id = result.get("id") if not job_id: raise HTTPException(status_code=500, detail="No job ID returned from RunPod") # Poll for completion for _ in range(120): # Max 10 minutes await asyncio.sleep(5) status_resp = await client.get( f"https://api.runpod.ai/v2/{RUNPOD_WHISPER_ENDPOINT}/status/{job_id}", headers={"Authorization": f"Bearer {RUNPOD_API_KEY}"}, ) status_data = status_resp.json() if status_data.get("status") == "COMPLETED": return status_data.get("output", {}) elif status_data.get("status") in ["FAILED", "CANCELLED"]: raise HTTPException( status_code=500, detail=f"Transcription failed: {status_data.get('error', 'Unknown error')}" ) raise HTTPException(status_code=504, detail="Transcription timed out") async def summarize_with_ollama(content: str, style: SummarizeStyle) -> str: """Summarize content using AI Orchestrator (Ollama)""" style_prompts = { SummarizeStyle.CONCISE: "Provide a concise 2-3 sentence summary of the following content:", SummarizeStyle.DETAILED: "Provide a detailed summary of the following content, covering all main points:", SummarizeStyle.BULLET_POINTS: "Summarize the following content as bullet points:", SummarizeStyle.TECHNICAL: "Provide a technical summary of the following content, focusing on key technical details:", SummarizeStyle.ELI5: "Explain the following content in simple terms that a child could understand:", } prompt = f"{style_prompts[style]}\n\n{content[:8000]}" # Limit content for context window async with httpx.AsyncClient(timeout=120) as client: try: resp = await client.post( f"{AI_ORCHESTRATOR_URL}/api/generate/text", json={ "prompt": prompt, "model": "llama3.2", "max_tokens": 1024, "priority": "low", # Use free Ollama }, ) result = resp.json() return result.get("response", "") except Exception as e: return f"[Summarization failed: {str(e)}]" async def index_to_semantic_search( doc_id: str, content: str, source: str, metadata: Dict[str, Any], ) -> bool: """Index document to semantic search service""" async with httpx.AsyncClient(timeout=30) as client: try: resp = await client.post( f"{SEMANTIC_SEARCH_URL}/index", json={ "id": doc_id, "content": content, "metadata": { "source": source, "indexed_at": datetime.now().isoformat(), **metadata, }, }, ) return resp.status_code == 200 except Exception: return False def extract_with_docling(file_path: Path, output_format: OutputFormat) -> tuple[str, dict]: """Extract content from document using Docling""" result = converter.convert(str(file_path)) doc = result.document # Get metadata metadata = { "pages": len(doc.pages) if hasattr(doc, "pages") else 0, "tables": len(doc.tables) if hasattr(doc, "tables") else 0, "figures": len(doc.pictures) if hasattr(doc, "pictures") else 0, } # Export in requested format if output_format == OutputFormat.MARKDOWN: content = doc.export_to_markdown() elif output_format == OutputFormat.JSON: content = doc.export_to_dict() elif output_format == OutputFormat.HTML: content = doc.export_to_html() else: # TEXT content = doc.export_to_markdown() # Markdown is readable as plain text return content if isinstance(content, str) else str(content), metadata # ============== API Endpoints ============== @app.get("/", response_class=HTMLResponse) async def dashboard(): """Web interface for document processing""" html = """ Document Intelligence - Extract, Summarize, Index

Document Intelligence

Extract, summarize, and index documents into your knowledge base

📄 Input Source

— or —
📁

Drop a file here or click to browse

PDF, DOCX, PPTX, XLSX, HTML, MD, TXT, EPUB, images, audio

⚙️ Processing Options

📊 Results

Processing...
-
Pages
-
Tables
-
Figures
-
Characters
Processing...
Loading content...
""" + str(stats["documents_processed"]) + """
Documents Processed
""" + str(stats["pages_extracted"]) + """
Pages Extracted
""" + str(stats["audio_transcribed"]) + """
Audio Transcribed
""" return HTMLResponse(content=html) @app.get("/api/info") async def api_info(): """Service info (JSON API endpoint)""" return { "service": "Docling Service", "version": "1.0.0", "status": "healthy", "supported_formats": { "documents": list(DOCUMENT_FORMATS), "images": list(IMAGE_FORMATS), "audio": list(AUDIO_FORMATS), }, "integrations": { "ai_orchestrator": AI_ORCHESTRATOR_URL, "semantic_search": SEMANTIC_SEARCH_URL, "runpod_whisper": f"endpoint:{RUNPOD_WHISPER_ENDPOINT}", }, } @app.get("/health") async def health(): """Health check endpoint for Traefik""" return {"status": "healthy", "timestamp": datetime.now().isoformat()} @app.get("/stats") async def get_stats(): """Get processing statistics""" return stats @app.post("/extract", response_model=ExtractionResult) async def extract_document(request: ExtractRequest): """ Extract content from a document (URL or base64). Supports: PDF, DOCX, PPTX, XLSX, HTML, MD, TXT, EPUB, images (with OCR) """ try: # Get content if request.url: content_bytes, filename = await fetch_url_content(str(request.url)) source = str(request.url) elif request.base64_content: content_bytes = base64.b64decode(request.base64_content) filename = request.filename or "document" source = f"base64:{filename}" else: raise HTTPException(status_code=400, detail="Provide either url or base64_content") ext = get_file_extension(filename) # Handle audio separately if ext in AUDIO_FORMATS: raise HTTPException( status_code=400, detail="Use /transcribe endpoint for audio files" ) # Write to temp file for Docling with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp: tmp.write(content_bytes) tmp_path = Path(tmp.name) try: # Extract content content, metadata = extract_with_docling(tmp_path, request.output_format) stats["documents_processed"] += 1 stats["pages_extracted"] += metadata.get("pages", 1) # Summarize if requested summary = None if request.summarize: summary = await summarize_with_ollama(content, request.summarize_style) # Index if requested indexed = False if request.index_to_search: doc_id = generate_doc_id(source, content) indexed = await index_to_semantic_search( doc_id=doc_id, content=content, source=source, metadata={**metadata, **(request.metadata or {})}, ) return ExtractionResult( success=True, source=source, content=content, format=request.output_format, metadata=metadata, summary=summary, indexed=indexed, ) finally: tmp_path.unlink(missing_ok=True) except HTTPException: raise except Exception as e: stats["errors"] += 1 return ExtractionResult( success=False, source=str(request.url or request.filename or "unknown"), format=request.output_format, error=str(e), ) @app.post("/extract/upload", response_model=ExtractionResult) async def extract_uploaded_file( file: UploadFile = File(...), output_format: OutputFormat = Form(OutputFormat.MARKDOWN), summarize: bool = Form(False), summarize_style: SummarizeStyle = Form(SummarizeStyle.CONCISE), index_to_search: bool = Form(False), ): """Extract content from an uploaded file""" try: content_bytes = await file.read() ext = get_file_extension(file.filename or "document") if ext in AUDIO_FORMATS: raise HTTPException( status_code=400, detail="Use /transcribe/upload endpoint for audio files" ) with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp: tmp.write(content_bytes) tmp_path = Path(tmp.name) try: content, metadata = extract_with_docling(tmp_path, output_format) stats["documents_processed"] += 1 stats["pages_extracted"] += metadata.get("pages", 1) summary = None if summarize: summary = await summarize_with_ollama(content, summarize_style) indexed = False if index_to_search: doc_id = generate_doc_id(file.filename or "upload", content) indexed = await index_to_semantic_search( doc_id=doc_id, content=content, source=f"upload:{file.filename}", metadata=metadata, ) return ExtractionResult( success=True, source=f"upload:{file.filename}", content=content, format=output_format, metadata=metadata, summary=summary, indexed=indexed, ) finally: tmp_path.unlink(missing_ok=True) except HTTPException: raise except Exception as e: stats["errors"] += 1 return ExtractionResult( success=False, source=f"upload:{file.filename}", format=output_format, error=str(e), ) @app.post("/transcribe") async def transcribe_audio(request: TranscribeRequest): """ Transcribe audio using RunPod Whisper. Supports: MP3, WAV, M4A, OGG, FLAC, WEBM """ try: if request.url: content_bytes, filename = await fetch_url_content(str(request.url)) source = str(request.url) elif request.base64_content: content_bytes = base64.b64decode(request.base64_content) source = "base64:audio" else: raise HTTPException(status_code=400, detail="Provide either url or base64_content") # Transcribe result = await transcribe_audio_runpod(content_bytes, request.language) stats["audio_transcribed"] += 1 transcript = result.get("transcription", result.get("text", "")) # Summarize if requested summary = None if request.summarize and transcript: summary = await summarize_with_ollama(transcript, request.summarize_style) return { "success": True, "source": source, "transcript": transcript, "language": result.get("detected_language"), "duration": result.get("duration"), "summary": summary, } except HTTPException: raise except Exception as e: stats["errors"] += 1 return { "success": False, "source": str(request.url or "base64"), "error": str(e), } @app.post("/transcribe/upload") async def transcribe_uploaded_audio( file: UploadFile = File(...), language: Optional[str] = Form(None), summarize: bool = Form(False), summarize_style: SummarizeStyle = Form(SummarizeStyle.CONCISE), ): """Transcribe uploaded audio file""" try: content_bytes = await file.read() result = await transcribe_audio_runpod(content_bytes, language) stats["audio_transcribed"] += 1 transcript = result.get("transcription", result.get("text", "")) summary = None if summarize and transcript: summary = await summarize_with_ollama(transcript, summarize_style) return { "success": True, "source": f"upload:{file.filename}", "transcript": transcript, "language": result.get("detected_language"), "duration": result.get("duration"), "summary": summary, } except HTTPException: raise except Exception as e: stats["errors"] += 1 return { "success": False, "source": f"upload:{file.filename}", "error": str(e), } @app.post("/batch") async def batch_extract(request: BatchExtractRequest, background_tasks: BackgroundTasks): """ Batch extract multiple documents. Returns immediately with job ID, processes in background. """ job_id = hashlib.sha256(str(datetime.now()).encode()).hexdigest()[:16] # For now, process synchronously (can be enhanced with Redis queue later) results = [] for item in request.items: result = await extract_document(item) results.append(result) return { "job_id": job_id, "total": len(request.items), "results": results, } @app.post("/url/preview") async def preview_url(url: HttpUrl): """Quick preview of URL content (first 500 chars of markdown)""" try: content_bytes, filename = await fetch_url_content(str(url)) stats["urls_fetched"] += 1 ext = get_file_extension(filename) with tempfile.NamedTemporaryFile(suffix=ext or ".html", delete=False) as tmp: tmp.write(content_bytes) tmp_path = Path(tmp.name) try: content, metadata = extract_with_docling(tmp_path, OutputFormat.MARKDOWN) return { "success": True, "url": str(url), "preview": content[:500] + ("..." if len(content) > 500 else ""), "full_length": len(content), "metadata": metadata, } finally: tmp_path.unlink(missing_ok=True) except Exception as e: return { "success": False, "url": str(url), "error": str(e), } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8081)