CoolFace
Apppublic

lamhieu/docsifer

sourceHugging Facemitupdated 4mo agoView on Hugging Face
13likes
main.py213 linesDownload Raw Back to docsifer
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