lamhieu/docsifer
13
1"""Application factory + ASGI entry point.2 3This module owns the FastAPI ``app`` instance, the lifespan-managed4singletons (``DocsiferService``, ``AnalyticsService``, safety primitives)5and the optional Gradio UI mount.6 7Run with::8 9 uvicorn docsifer.main:app --host 0.0.0.0 --port 7860 --proxy-headers10"""11 12from __future__ import annotations13 14import asyncio15import contextlib16import logging17from collections.abc import AsyncIterator18from contextlib import asynccontextmanager19 20from fastapi import FastAPI21from fastapi.middleware.cors import CORSMiddleware22from fastapi.middleware.gzip import GZipMiddleware23from fastapi.responses import ORJSONResponse24 25from .analytics import AnalyticsService, InMemoryStore, UpstashStore26from .analytics.store import AnalyticsStore27from .api.error_handlers import register_exception_handlers28from .api.middleware import register_middleware29from .api.v1 import router as v1_router30from .config import Settings, get_settings31from .core.service import DocsiferService32from .logging_config import configure_logging33from .safety import (34 ConversionGate,35 PerIPLimiter,36 ResourceGuard,37 disk_cleanup_loop,38 memory_watchdog_loop,39)40 41logger = logging.getLogger(__name__)42 43 44def _build_analytics_store(settings: Settings) -> AnalyticsStore:45 if not settings.analytics_persistent:46 logger.info("Analytics: in-memory (no DOCSIFER_REDIS_URL configured)")47 return InMemoryStore()48 try:49 store = UpstashStore(url=settings.redis_url, token=settings.redis_token)50 logger.info("Analytics: Upstash store configured")51 return store52 except Exception as exc: # pragma: no cover - defensive53 logger.warning("Could not initialize Upstash store (%s); falling back to in-memory", exc)54 return InMemoryStore()55 56 57@asynccontextmanager58async def _lifespan(app: FastAPI) -> AsyncIterator[None]:59 settings: Settings = app.state.settings60 configure_logging(level=settings.log_level, json_output=settings.log_json)61 logger.info(62 "Starting %s v%s (env=%s)",63 settings.app_name,64 settings.app_version,65 settings.environment,66 )67 68 # ---- core converter -------------------------------------------------69 converter = DocsiferService(70 token_model=settings.token_model,71 default_openai_base_url=settings.default_openai_base_url,72 default_openai_model=settings.default_openai_model,73 worker_pool_size=settings.worker_pool_size,74 llm_cache_max_size=settings.llm_cache_max_size,75 llm_cache_ttl=settings.llm_cache_ttl_sec,76 openai_request_timeout=settings.openai_request_timeout_sec,77 openai_connect_timeout=settings.openai_connect_timeout_sec,78 openai_max_retries=settings.openai_max_retries,79 known_extensions=set(settings.allowed_extensions),80 )81 app.state.converter = converter82 83 # ---- analytics ------------------------------------------------------84 store = _build_analytics_store(settings)85 analytics = AnalyticsService(86 store=store,87 sync_interval_sec=settings.analytics_sync_interval_sec,88 max_retries=settings.analytics_max_retries,89 label=settings.analytics_label,90 )91 await analytics.start()92 app.state.analytics = analytics93 94 # ---- safety primitives ---------------------------------------------95 app.state.conversion_gate = ConversionGate(96 max_concurrent=settings.max_concurrent_conversions,97 max_queue=settings.max_queue_depth,98 )99 app.state.per_ip_limiter = PerIPLimiter(100 max_per_ip=settings.max_per_ip_concurrent,101 )102 app.state.resource_guard = ResourceGuard(103 min_free_memory_mb=settings.min_free_memory_mb,104 min_free_disk_mb=settings.min_free_disk_mb,105 tmp_dir=settings.tmp_dir,106 )107 108 # ---- background loops ----------------------------------------------109 stop_event = asyncio.Event()110 background_tasks: list[asyncio.Task[None]] = [111 asyncio.create_task(112 disk_cleanup_loop(113 tmp_dir=settings.tmp_dir,114 ttl_sec=settings.disk_cleanup_ttl_sec,115 interval_sec=settings.disk_cleanup_interval_sec,116 stop_event=stop_event,117 ),118 name="docsifer-disk-cleanup",119 )120 ]121 if settings.enable_memory_watchdog:122 background_tasks.append(123 asyncio.create_task(124 memory_watchdog_loop(125 threshold_pct=settings.memory_watchdog_pct,126 interval_sec=settings.memory_watchdog_interval_sec,127 stop_event=stop_event,128 ),129 name="docsifer-memory-watchdog",130 )131 )132 app.state.stop_event = stop_event133 134 try:135 yield136 finally:137 logger.info("Shutting down %s", settings.app_name)138 stop_event.set()139 for task in background_tasks:140 task.cancel()141 with contextlib.suppress(asyncio.CancelledError, Exception):142 await task143 with contextlib.suppress(Exception):144 await analytics.stop()145 with contextlib.suppress(Exception):146 await converter.shutdown()147 148 149def create_app(settings: Settings | None = None) -> FastAPI:150 """Build and return the FastAPI application."""151 settings = settings or get_settings()152 153 app = FastAPI(154 title=f"{settings.app_name} Service API",155 description=(156 "Convert PDF, PowerPoint, Word, Excel, images, audio, HTML, JSON, "157 "CSV, XML, ZIP, and more into Markdown — optionally enhanced by an LLM."158 ),159 version=settings.app_version,160 docs_url="/docs",161 redoc_url="/redoc",162 default_response_class=ORJSONResponse,163 lifespan=_lifespan,164 )165 app.state.settings = settings166 167 # CORS — fixes Bug A9 (allow_credentials with origins='*')168 app.add_middleware(169 CORSMiddleware,170 allow_origins=settings.cors_origins,171 allow_credentials=settings.cors_allow_credentials_safe,172 allow_methods=["*"],173 allow_headers=["*"],174 expose_headers=["X-Request-ID"],175 )176 # GZip large responses (Section B.4)177 app.add_middleware(GZipMiddleware, minimum_size=settings.gzip_min_size)178 179 register_middleware(app, settings)180 register_exception_handlers(app)181 182 # ---- routes ---------------------------------------------------------183 app.include_router(v1_router, prefix="/v1")184 185 # ---- optional Gradio UI --------------------------------------------186 _maybe_mount_gradio(app, settings)187 188 return app189 190 191def _maybe_mount_gradio(app: FastAPI, settings: Settings) -> None:192 try:193 import gradio as gr # type: ignore194 195 mount_fn = getattr(gr, "mount_gradio_app", None)196 if mount_fn is None:197 from gradio.routes import mount_gradio_app as mount_fn # type: ignore198 from .ui.gradio_app import build_interface199 except Exception as exc: # pragma: no cover - optional dependency200 logger.warning("Gradio UI disabled (%s)", exc)201 return202 203 try:204 interface = build_interface(settings, app)205 mount_fn(app, interface, path="/")206 logger.info("Gradio UI mounted at /")207 except Exception as exc: # pragma: no cover - defensive208 logger.warning("Could not mount Gradio UI (%s)", exc)209 210 211# Public ASGI app instance ----------------------------------------------------212app: FastAPI = create_app()213 