Merge pull request #41 from scub-france/feature/quality-improvements
Feature/quality improvements
This commit is contained in:
Pier-Jean Malandrino
GitHub
commit
4e6fa0908d
20 changed files with 654 additions and 62 deletions
|
|
@ -46,7 +46,7 @@ def _to_response(job) -> AnalysisResponse:
|
|||
|
||||
|
||||
@router.post("", response_model=AnalysisResponse)
|
||||
async def create_analysis(body: CreateAnalysisRequest, service: ServiceDep):
|
||||
async def create_analysis(body: CreateAnalysisRequest, service: ServiceDep) -> AnalysisResponse:
|
||||
"""Create a new analysis job for a document."""
|
||||
if not body.documentId or not body.documentId.strip():
|
||||
raise HTTPException(status_code=400, detail="documentId is required")
|
||||
|
|
@ -72,14 +72,14 @@ async def create_analysis(body: CreateAnalysisRequest, service: ServiceDep):
|
|||
|
||||
|
||||
@router.get("", response_model=list[AnalysisResponse])
|
||||
async def list_analyses(service: ServiceDep):
|
||||
async def list_analyses(service: ServiceDep) -> list[AnalysisResponse]:
|
||||
"""List all analysis jobs."""
|
||||
jobs = await service.find_all()
|
||||
return [_to_response(j) for j in jobs]
|
||||
|
||||
|
||||
@router.get("/{job_id}", response_model=AnalysisResponse)
|
||||
async def get_analysis(job_id: str, service: ServiceDep):
|
||||
async def get_analysis(job_id: str, service: ServiceDep) -> AnalysisResponse:
|
||||
"""Get a single analysis job."""
|
||||
job = await service.find_by_id(job_id)
|
||||
if not job:
|
||||
|
|
@ -88,7 +88,9 @@ async def get_analysis(job_id: str, service: ServiceDep):
|
|||
|
||||
|
||||
@router.post("/{job_id}/rechunk", response_model=list[ChunkResponse])
|
||||
async def rechunk_analysis(job_id: str, body: RechunkRequest, service: ServiceDep):
|
||||
async def rechunk_analysis(
|
||||
job_id: str, body: RechunkRequest, service: ServiceDep
|
||||
) -> list[ChunkResponse]:
|
||||
"""Re-chunk a completed analysis with new chunking options."""
|
||||
try:
|
||||
chunks = await service.rechunk(job_id, body.chunkingOptions.model_dump())
|
||||
|
|
@ -107,7 +109,7 @@ async def rechunk_analysis(job_id: str, body: RechunkRequest, service: ServiceDe
|
|||
|
||||
|
||||
@router.delete("/{job_id}", status_code=204)
|
||||
async def delete_analysis(job_id: str, service: ServiceDep):
|
||||
async def delete_analysis(job_id: str, service: ServiceDep) -> None:
|
||||
"""Delete an analysis job."""
|
||||
deleted = await service.delete(job_id)
|
||||
if not deleted:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ from services import document_service
|
|||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/documents", tags=["documents"])
|
||||
|
||||
_READ_CHUNK_SIZE = 64 * 1024 # 64 KB
|
||||
|
||||
|
||||
def _to_response(doc) -> DocumentResponse:
|
||||
return DocumentResponse(
|
||||
|
|
@ -25,8 +27,8 @@ def _to_response(doc) -> DocumentResponse:
|
|||
)
|
||||
|
||||
|
||||
@router.post("/upload", response_model=DocumentResponse)
|
||||
async def upload(file: UploadFile):
|
||||
@router.post("/upload", response_model=DocumentResponse, status_code=200)
|
||||
async def upload(file: UploadFile) -> DocumentResponse:
|
||||
"""Upload a PDF document."""
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="No filename provided")
|
||||
|
|
@ -35,7 +37,15 @@ async def upload(file: UploadFile):
|
|||
if file.size and file.size > document_service.MAX_FILE_SIZE:
|
||||
raise HTTPException(status_code=413, detail="File too large (max 50 MB)")
|
||||
|
||||
content = await file.read()
|
||||
# Read in chunks to avoid holding the full upload in a single allocation
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
while chunk := await file.read(_READ_CHUNK_SIZE):
|
||||
total += len(chunk)
|
||||
if total > document_service.MAX_FILE_SIZE:
|
||||
raise HTTPException(status_code=413, detail="File too large (max 50 MB)")
|
||||
chunks.append(chunk)
|
||||
content = b"".join(chunks)
|
||||
|
||||
try:
|
||||
doc = await document_service.upload(
|
||||
|
|
@ -44,20 +54,20 @@ async def upload(file: UploadFile):
|
|||
file_content=content,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=413, detail=str(e)) from e
|
||||
raise HTTPException(status_code=400, detail=str(e)) from e
|
||||
|
||||
return _to_response(doc)
|
||||
|
||||
|
||||
@router.get("", response_model=list[DocumentResponse])
|
||||
async def list_documents():
|
||||
async def list_documents() -> list[DocumentResponse]:
|
||||
"""List all documents."""
|
||||
docs = await document_service.find_all()
|
||||
return [_to_response(d) for d in docs]
|
||||
|
||||
|
||||
@router.get("/{doc_id}", response_model=DocumentResponse)
|
||||
async def get_document(doc_id: str):
|
||||
async def get_document(doc_id: str) -> DocumentResponse:
|
||||
"""Get a single document."""
|
||||
doc = await document_service.find_by_id(doc_id)
|
||||
if not doc:
|
||||
|
|
@ -66,7 +76,7 @@ async def get_document(doc_id: str):
|
|||
|
||||
|
||||
@router.delete("/{doc_id}", status_code=204)
|
||||
async def delete_document(doc_id: str):
|
||||
async def delete_document(doc_id: str) -> None:
|
||||
"""Delete a document and its file."""
|
||||
deleted = await document_service.delete(doc_id)
|
||||
if not deleted:
|
||||
|
|
@ -78,12 +88,18 @@ async def preview(
|
|||
doc_id: str,
|
||||
page: int = Query(1, ge=1),
|
||||
dpi: int = Query(150, ge=72, le=300),
|
||||
):
|
||||
) -> Response:
|
||||
"""Generate a PNG preview of a specific PDF page."""
|
||||
doc = await document_service.find_by_id(doc_id)
|
||||
if not doc:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
|
||||
if doc.page_count and page > doc.page_count:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Page {page} out of range (document has {doc.page_count} pages)",
|
||||
)
|
||||
|
||||
try:
|
||||
with open(doc.storage_path, "rb") as f:
|
||||
file_content = f.read()
|
||||
|
|
@ -91,6 +107,11 @@ async def preview(
|
|||
return Response(content=png_bytes, media_type="image/png")
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e)) from e
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail="PDF file not found on disk") from exc
|
||||
except OSError as exc:
|
||||
logger.exception("I/O error generating preview for %s", doc_id)
|
||||
raise HTTPException(status_code=422, detail="Failed to read PDF file") from exc
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to generate preview")
|
||||
logger.exception("Unexpected error generating preview for %s", doc_id)
|
||||
raise HTTPException(status_code=422, detail="Failed to generate preview") from exc
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from __future__ import annotations
|
|||
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, field_validator
|
||||
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
|
||||
def _to_camel(name: str) -> str:
|
||||
|
|
@ -29,6 +29,7 @@ class _CamelModel(BaseModel):
|
|||
class DocumentResponse(_CamelModel):
|
||||
id: str
|
||||
filename: str
|
||||
status: str = "uploaded" # Document status (always "uploaded" for now)
|
||||
content_type: str | None = None
|
||||
file_size: int | None = None
|
||||
page_count: int | None = None
|
||||
|
|
@ -54,16 +55,39 @@ class AnalysisResponse(_CamelModel):
|
|||
class PipelineOptionsRequest(BaseModel):
|
||||
"""Docling pipeline configuration options."""
|
||||
|
||||
do_ocr: bool = True
|
||||
do_table_structure: bool = True
|
||||
table_mode: str = "accurate" # "accurate" or "fast"
|
||||
do_code_enrichment: bool = False
|
||||
do_formula_enrichment: bool = False
|
||||
do_picture_classification: bool = False
|
||||
do_picture_description: bool = False
|
||||
generate_picture_images: bool = False
|
||||
generate_page_images: bool = False
|
||||
images_scale: float = 1.0
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
do_ocr: bool = Field(default=True, validation_alias=AliasChoices("do_ocr", "doOcr"))
|
||||
do_table_structure: bool = Field(
|
||||
default=True, validation_alias=AliasChoices("do_table_structure", "doTableStructure")
|
||||
)
|
||||
table_mode: str = Field(
|
||||
default="accurate", validation_alias=AliasChoices("table_mode", "tableMode")
|
||||
)
|
||||
do_code_enrichment: bool = Field(
|
||||
default=False, validation_alias=AliasChoices("do_code_enrichment", "doCodeEnrichment")
|
||||
)
|
||||
do_formula_enrichment: bool = Field(
|
||||
default=False, validation_alias=AliasChoices("do_formula_enrichment", "doFormulaEnrichment")
|
||||
)
|
||||
do_picture_classification: bool = Field(
|
||||
default=False,
|
||||
validation_alias=AliasChoices("do_picture_classification", "doPictureClassification"),
|
||||
)
|
||||
do_picture_description: bool = Field(
|
||||
default=False,
|
||||
validation_alias=AliasChoices("do_picture_description", "doPictureDescription"),
|
||||
)
|
||||
generate_picture_images: bool = Field(
|
||||
default=False,
|
||||
validation_alias=AliasChoices("generate_picture_images", "generatePictureImages"),
|
||||
)
|
||||
generate_page_images: bool = Field(
|
||||
default=False, validation_alias=AliasChoices("generate_page_images", "generatePageImages")
|
||||
)
|
||||
images_scale: float = Field(
|
||||
default=1.0, validation_alias=AliasChoices("images_scale", "imagesScale")
|
||||
)
|
||||
|
||||
@field_validator("table_mode")
|
||||
@classmethod
|
||||
|
|
@ -83,10 +107,18 @@ class PipelineOptionsRequest(BaseModel):
|
|||
class ChunkingOptionsRequest(BaseModel):
|
||||
"""Docling chunking configuration options."""
|
||||
|
||||
chunker_type: str = "hybrid" # "hybrid", "hierarchical"
|
||||
max_tokens: int = 512
|
||||
merge_peers: bool = True
|
||||
repeat_table_header: bool = True
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
chunker_type: str = Field(
|
||||
default="hybrid", validation_alias=AliasChoices("chunker_type", "chunkerType")
|
||||
)
|
||||
max_tokens: int = Field(default=512, validation_alias=AliasChoices("max_tokens", "maxTokens"))
|
||||
merge_peers: bool = Field(
|
||||
default=True, validation_alias=AliasChoices("merge_peers", "mergePeers")
|
||||
)
|
||||
repeat_table_header: bool = Field(
|
||||
default=True, validation_alias=AliasChoices("repeat_table_header", "repeatTableHeader")
|
||||
)
|
||||
|
||||
@field_validator("chunker_type")
|
||||
@classmethod
|
||||
|
|
@ -117,10 +149,16 @@ class ChunkResponse(_CamelModel):
|
|||
|
||||
|
||||
class CreateAnalysisRequest(BaseModel):
|
||||
documentId: str # camelCase to match existing frontend contract
|
||||
pipelineOptions: PipelineOptionsRequest | None = None
|
||||
chunkingOptions: ChunkingOptionsRequest | None = None
|
||||
documentId: str = Field(validation_alias=AliasChoices("documentId", "document_id"))
|
||||
pipelineOptions: PipelineOptionsRequest | None = Field(
|
||||
default=None, validation_alias=AliasChoices("pipelineOptions", "pipeline_options")
|
||||
)
|
||||
chunkingOptions: ChunkingOptionsRequest | None = Field(
|
||||
default=None, validation_alias=AliasChoices("chunkingOptions", "chunking_options")
|
||||
)
|
||||
|
||||
|
||||
class RechunkRequest(BaseModel):
|
||||
chunkingOptions: ChunkingOptionsRequest
|
||||
chunkingOptions: ChunkingOptionsRequest = Field(
|
||||
validation_alias=AliasChoices("chunkingOptions", "chunking_options")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ class AnalysisJob:
|
|||
document_filename: str | None = None
|
||||
|
||||
def mark_running(self) -> None:
|
||||
"""Transition to RUNNING and record the start timestamp."""
|
||||
self.status = AnalysisStatus.RUNNING
|
||||
self.started_at = _utcnow()
|
||||
|
||||
|
|
@ -64,6 +65,7 @@ class AnalysisJob:
|
|||
document_json: str | None = None,
|
||||
chunks_json: str | None = None,
|
||||
) -> None:
|
||||
"""Transition to COMPLETED with conversion results."""
|
||||
self.status = AnalysisStatus.COMPLETED
|
||||
self.content_markdown = markdown
|
||||
self.content_html = html
|
||||
|
|
@ -73,6 +75,7 @@ class AnalysisJob:
|
|||
self.completed_at = _utcnow()
|
||||
|
||||
def mark_failed(self, error: str) -> None:
|
||||
"""Transition to FAILED with an error message."""
|
||||
self.status = AnalysisStatus.FAILED
|
||||
self.error_message = error
|
||||
self.completed_at = _utcnow()
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ class ConversionOptions:
|
|||
images_scale: float = 1.0
|
||||
|
||||
def is_default(self) -> bool:
|
||||
"""Return True if all options match their defaults."""
|
||||
return self == ConversionOptions()
|
||||
|
||||
|
||||
|
|
@ -60,6 +61,7 @@ class ChunkingOptions:
|
|||
repeat_table_header: bool = True
|
||||
|
||||
def is_default(self) -> bool:
|
||||
"""Return True if all options match their defaults."""
|
||||
return self == ChunkingOptions()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ from docling_core.transforms.chunker import HierarchicalChunker
|
|||
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from infra.bbox import EMPTY_BBOX, to_topleft_list
|
||||
from domain.value_objects import ChunkBbox, ChunkingOptions, ChunkResult
|
||||
from infra.bbox import EMPTY_BBOX, to_topleft_list
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
|
|||
|
|
@ -35,13 +35,13 @@ from docling_core.types.doc import (
|
|||
TitleItem,
|
||||
)
|
||||
|
||||
from infra.bbox import to_topleft_list
|
||||
from domain.value_objects import (
|
||||
ConversionOptions,
|
||||
ConversionResult,
|
||||
PageDetail,
|
||||
PageElement,
|
||||
)
|
||||
from infra.bbox import to_topleft_list
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
|
|||
89
document-parser/infra/rate_limiter.py
Normal file
89
document-parser/infra/rate_limiter.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
"""Lightweight in-memory rate limiter middleware for FastAPI.
|
||||
|
||||
Uses a sliding-window counter per client IP. No external dependency
|
||||
required — suitable for single-process deployments with SQLite.
|
||||
|
||||
For multi-process or distributed setups, replace with a Redis-backed
|
||||
solution (e.g. slowapi).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||
from starlette.responses import JSONResponse, Response
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ClientBucket:
|
||||
"""Sliding window of request timestamps for a single client."""
|
||||
|
||||
timestamps: list[float] = field(default_factory=list)
|
||||
|
||||
def count_recent(self, window: float, now: float) -> int:
|
||||
"""Remove expired entries and return the count of recent requests."""
|
||||
cutoff = now - window
|
||||
self.timestamps = [t for t in self.timestamps if t > cutoff]
|
||||
return len(self.timestamps)
|
||||
|
||||
def add(self, now: float) -> None:
|
||||
self.timestamps.append(now)
|
||||
|
||||
|
||||
class RateLimiterMiddleware(BaseHTTPMiddleware):
|
||||
"""Per-IP rate limiter using in-memory sliding windows.
|
||||
|
||||
Args:
|
||||
app: The ASGI application.
|
||||
requests_per_window: Max requests allowed per window.
|
||||
window_seconds: Size of the sliding window in seconds.
|
||||
exclude_paths: Paths exempt from rate limiting (e.g. health checks).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app,
|
||||
*,
|
||||
requests_per_window: int = 60,
|
||||
window_seconds: float = 60.0,
|
||||
exclude_paths: tuple[str, ...] = ("/api/health",),
|
||||
):
|
||||
super().__init__(app)
|
||||
self._max_requests = requests_per_window
|
||||
self._window = window_seconds
|
||||
self._exclude = exclude_paths
|
||||
self._buckets: dict[str, _ClientBucket] = defaultdict(_ClientBucket)
|
||||
|
||||
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||
if request.url.path in self._exclude:
|
||||
return await call_next(request)
|
||||
|
||||
client_ip = request.client.host if request.client else "unknown"
|
||||
now = time.monotonic()
|
||||
|
||||
bucket = self._buckets[client_ip]
|
||||
recent = bucket.count_recent(self._window, now)
|
||||
|
||||
if recent >= self._max_requests:
|
||||
retry_after = int(self._window)
|
||||
logger.warning(
|
||||
"Rate limit exceeded for %s (%d/%d)", client_ip, recent, self._max_requests
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=429,
|
||||
content={"detail": "Too many requests"},
|
||||
headers={"Retry-After": str(retry_after)},
|
||||
)
|
||||
|
||||
bucket.add(now)
|
||||
return await call_next(request)
|
||||
|
|
@ -13,6 +13,7 @@ class Settings:
|
|||
docling_serve_url: str = "http://localhost:5001"
|
||||
docling_serve_api_key: str | None = None
|
||||
conversion_timeout: int = 600
|
||||
max_concurrent_analyses: int = 3
|
||||
upload_dir: str = "./uploads"
|
||||
db_path: str = "./data/docling_studio.db"
|
||||
cors_origins: list[str] = field(
|
||||
|
|
@ -21,6 +22,7 @@ class Settings:
|
|||
|
||||
@classmethod
|
||||
def from_env(cls) -> Settings:
|
||||
"""Build a Settings instance from environment variables."""
|
||||
cors_raw = os.environ.get("CORS_ORIGINS", "http://localhost:3000,http://localhost:5173")
|
||||
return cls(
|
||||
app_version=os.environ.get("APP_VERSION", "dev"),
|
||||
|
|
@ -28,7 +30,12 @@ class Settings:
|
|||
docling_serve_url=os.environ.get("DOCLING_SERVE_URL", "http://localhost:5001"),
|
||||
docling_serve_api_key=os.environ.get("DOCLING_SERVE_API_KEY"),
|
||||
conversion_timeout=int(os.environ.get("CONVERSION_TIMEOUT", "600")),
|
||||
max_concurrent_analyses=int(os.environ.get("MAX_CONCURRENT_ANALYSES", "3")),
|
||||
upload_dir=os.environ.get("UPLOAD_DIR", "./uploads"),
|
||||
db_path=os.environ.get("DB_PATH", "./data/docling_studio.db"),
|
||||
cors_origins=[o.strip() for o in cors_raw.split(",")],
|
||||
)
|
||||
|
||||
|
||||
# Module-level singleton — import this from other modules.
|
||||
settings = Settings.from_env()
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ Conversion engine is selected via CONVERSION_ENGINE env var:
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
|
@ -19,8 +20,9 @@ from fastapi.middleware.cors import CORSMiddleware
|
|||
|
||||
from api.analyses import router as analyses_router
|
||||
from api.documents import router as documents_router
|
||||
from infra.settings import Settings
|
||||
from persistence.database import init_db
|
||||
from infra.rate_limiter import RateLimiterMiddleware
|
||||
from infra.settings import settings
|
||||
from persistence.database import get_connection, init_db
|
||||
from services.analysis_service import AnalysisService
|
||||
|
||||
logging.basicConfig(
|
||||
|
|
@ -29,12 +31,6 @@ logging.basicConfig(
|
|||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Settings & dependency wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
settings = Settings.from_env()
|
||||
|
||||
|
||||
def _build_converter():
|
||||
"""Build the converter adapter based on configuration."""
|
||||
|
|
@ -69,6 +65,7 @@ def _build_analysis_service() -> AnalysisService:
|
|||
converter=converter,
|
||||
chunker=chunker,
|
||||
conversion_timeout=settings.conversion_timeout,
|
||||
max_concurrent=settings.max_concurrent_analyses,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -78,7 +75,7 @@ def _build_analysis_service() -> AnalysisService:
|
|||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
await init_db()
|
||||
app.state.analysis_service = _build_analysis_service()
|
||||
logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine)
|
||||
|
|
@ -98,16 +95,27 @@ app.add_middleware(
|
|||
allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
|
||||
allow_headers=["Content-Type", "Authorization"],
|
||||
)
|
||||
app.add_middleware(RateLimiterMiddleware, requests_per_window=100, window_seconds=60)
|
||||
|
||||
app.include_router(documents_router)
|
||||
app.include_router(analyses_router)
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
def health():
|
||||
"""Health check endpoint."""
|
||||
async def health() -> dict[str, str]:
|
||||
"""Health check endpoint — verifies database connectivity."""
|
||||
db_status = "ok"
|
||||
try:
|
||||
async with get_connection() as db:
|
||||
await db.execute("SELECT 1")
|
||||
except Exception:
|
||||
db_status = "error"
|
||||
logger.warning("Health check: database unreachable", exc_info=True)
|
||||
|
||||
status = "ok" if db_status == "ok" else "degraded"
|
||||
return {
|
||||
"status": "ok",
|
||||
"status": status,
|
||||
"version": settings.app_version,
|
||||
"engine": settings.conversion_engine,
|
||||
"database": db_status,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ _SELECT_WITH_DOC = """
|
|||
|
||||
|
||||
async def insert(job: AnalysisJob) -> None:
|
||||
"""Persist a new analysis job record."""
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"""INSERT INTO analysis_jobs (id, document_id, status, created_at)
|
||||
|
|
@ -52,6 +53,7 @@ async def insert(job: AnalysisJob) -> None:
|
|||
|
||||
|
||||
async def find_all(*, limit: int = 200, offset: int = 0) -> list[AnalysisJob]:
|
||||
"""Return analysis jobs with document info, newest first."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute(
|
||||
f"{_SELECT_WITH_DOC} ORDER BY aj.created_at DESC LIMIT ? OFFSET ?",
|
||||
|
|
@ -62,6 +64,7 @@ async def find_all(*, limit: int = 200, offset: int = 0) -> list[AnalysisJob]:
|
|||
|
||||
|
||||
async def find_by_id(job_id: str) -> AnalysisJob | None:
|
||||
"""Find an analysis job by ID (with document filename), or return None."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute(f"{_SELECT_WITH_DOC} WHERE aj.id = ?", (job_id,))
|
||||
row = await cursor.fetchone()
|
||||
|
|
@ -69,6 +72,7 @@ async def find_by_id(job_id: str) -> AnalysisJob | None:
|
|||
|
||||
|
||||
async def update_status(job: AnalysisJob) -> None:
|
||||
"""Persist all mutable fields of an analysis job (status, results, timestamps)."""
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"""UPDATE analysis_jobs
|
||||
|
|
@ -104,6 +108,7 @@ async def update_chunks(job_id: str, chunks_json: str) -> bool:
|
|||
|
||||
|
||||
async def delete(job_id: str) -> bool:
|
||||
"""Delete an analysis job by ID. Returns True if a row was removed."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute("DELETE FROM analysis_jobs WHERE id = ?", (job_id,))
|
||||
await db.commit()
|
||||
|
|
|
|||
|
|
@ -4,13 +4,16 @@ from __future__ import annotations
|
|||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from infra.settings import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DB_PATH = os.environ.get("DB_PATH", "./data/docling_studio.db")
|
||||
DB_PATH = settings.db_path
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
|
|
@ -80,7 +83,7 @@ async def get_db() -> aiosqlite.Connection:
|
|||
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_connection():
|
||||
async def get_connection() -> AsyncIterator[aiosqlite.Connection]:
|
||||
"""Context manager that opens and auto-closes a database connection."""
|
||||
db = await get_db()
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ def _row_to_document(row) -> Document:
|
|||
|
||||
|
||||
async def insert(doc: Document) -> None:
|
||||
"""Persist a new document record."""
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"""INSERT INTO documents (id, filename, content_type, file_size, page_count, storage_path, created_at)
|
||||
|
|
@ -44,6 +45,7 @@ async def insert(doc: Document) -> None:
|
|||
|
||||
|
||||
async def find_all(*, limit: int = 200, offset: int = 0) -> list[Document]:
|
||||
"""Return documents ordered by creation date (newest first)."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute(
|
||||
"SELECT * FROM documents ORDER BY created_at DESC LIMIT ? OFFSET ?",
|
||||
|
|
@ -54,6 +56,7 @@ async def find_all(*, limit: int = 200, offset: int = 0) -> list[Document]:
|
|||
|
||||
|
||||
async def find_by_id(doc_id: str) -> Document | None:
|
||||
"""Find a document by its ID, or return None."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute("SELECT * FROM documents WHERE id = ?", (doc_id,))
|
||||
row = await cursor.fetchone()
|
||||
|
|
@ -61,6 +64,7 @@ async def find_by_id(doc_id: str) -> Document | None:
|
|||
|
||||
|
||||
async def update_page_count(doc_id: str, page_count: int) -> None:
|
||||
"""Update the page count after conversion has determined it."""
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"UPDATE documents SET page_count = ? WHERE id = ?",
|
||||
|
|
@ -70,6 +74,7 @@ async def update_page_count(doc_id: str, page_count: int) -> None:
|
|||
|
||||
|
||||
async def delete(doc_id: str) -> bool:
|
||||
"""Delete a document by ID. Returns True if a row was removed."""
|
||||
async with get_connection() as db:
|
||||
cursor = await db.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
|
||||
await db.commit()
|
||||
|
|
|
|||
|
|
@ -34,6 +34,10 @@ def _chunk_to_dict(c: ChunkResult) -> dict:
|
|||
}
|
||||
|
||||
|
||||
# Maximum number of concurrent analysis jobs to prevent resource exhaustion.
|
||||
_DEFAULT_MAX_CONCURRENT = 3
|
||||
|
||||
|
||||
class AnalysisService:
|
||||
"""Orchestrates document analysis using an injected converter."""
|
||||
|
||||
|
|
@ -42,10 +46,12 @@ class AnalysisService:
|
|||
converter: DocumentConverter,
|
||||
chunker: DocumentChunker | None = None,
|
||||
conversion_timeout: int = 600,
|
||||
max_concurrent: int = _DEFAULT_MAX_CONCURRENT,
|
||||
):
|
||||
self._converter = converter
|
||||
self._chunker = chunker
|
||||
self._conversion_timeout = conversion_timeout
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
async def create(
|
||||
self,
|
||||
|
|
@ -77,12 +83,15 @@ class AnalysisService:
|
|||
return job
|
||||
|
||||
async def find_all(self) -> list[AnalysisJob]:
|
||||
"""Return all analysis jobs, newest first."""
|
||||
return await analysis_repo.find_all()
|
||||
|
||||
async def find_by_id(self, job_id: str) -> AnalysisJob | None:
|
||||
"""Find an analysis job by ID, or return None."""
|
||||
return await analysis_repo.find_by_id(job_id)
|
||||
|
||||
async def delete(self, job_id: str) -> bool:
|
||||
"""Delete an analysis job. Returns True if it existed."""
|
||||
return await analysis_repo.delete(job_id)
|
||||
|
||||
async def rechunk(self, job_id: str, chunking_options: dict) -> list[ChunkResult]:
|
||||
|
|
@ -113,7 +122,25 @@ class AnalysisService:
|
|||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> None:
|
||||
"""Background task: run conversion and optionally chunk."""
|
||||
"""Background task: run conversion and optionally chunk.
|
||||
|
||||
Acquires the concurrency semaphore to limit parallel conversions
|
||||
and prevent CPU/memory exhaustion on modest hardware.
|
||||
"""
|
||||
async with self._semaphore:
|
||||
await self._run_analysis_inner(
|
||||
job_id, file_path, filename, pipeline_options, chunking_options
|
||||
)
|
||||
|
||||
async def _run_analysis_inner(
|
||||
self,
|
||||
job_id: str,
|
||||
file_path: str,
|
||||
filename: str,
|
||||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> None:
|
||||
"""Inner analysis logic — called under the concurrency semaphore."""
|
||||
try:
|
||||
job = await analysis_repo.find_by_id(job_id)
|
||||
if not job:
|
||||
|
|
@ -192,5 +219,7 @@ async def _mark_failed(job_id: str, error: str) -> None:
|
|||
if job:
|
||||
job.mark_failed(error)
|
||||
await analysis_repo.update_status(job)
|
||||
except OSError:
|
||||
logger.exception("Database I/O error marking job %s as failed", job_id)
|
||||
except Exception:
|
||||
logger.exception("Could not mark job %s as failed", job_id)
|
||||
logger.exception("Unexpected error marking job %s as failed", job_id)
|
||||
|
|
|
|||
|
|
@ -10,11 +10,12 @@ import uuid
|
|||
from pdf2image import convert_from_bytes, pdfinfo_from_bytes
|
||||
|
||||
from domain.models import Document
|
||||
from infra.settings import settings
|
||||
from persistence import analysis_repo, document_repo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
UPLOAD_DIR = os.environ.get("UPLOAD_DIR", "./uploads")
|
||||
UPLOAD_DIR = settings.upload_dir
|
||||
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB
|
||||
|
||||
|
||||
|
|
@ -22,8 +23,14 @@ MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB
|
|||
_PDF_MAGIC = b"%PDF"
|
||||
|
||||
|
||||
_UPLOAD_CHUNK_SIZE = 64 * 1024 # 64 KB chunks for streaming writes
|
||||
|
||||
|
||||
async def upload(filename: str, content_type: str, file_content: bytes) -> Document:
|
||||
"""Save uploaded file to disk and persist metadata."""
|
||||
"""Save uploaded file to disk and persist metadata.
|
||||
|
||||
Writes the file in fixed-size chunks to keep peak memory usage low.
|
||||
"""
|
||||
if len(file_content) > MAX_FILE_SIZE:
|
||||
raise ValueError("File too large (max 50 MB)")
|
||||
|
||||
|
|
@ -36,8 +43,10 @@ async def upload(filename: str, content_type: str, file_content: bytes) -> Docum
|
|||
safe_name = f"{uuid.uuid4()}{ext}"
|
||||
file_path = os.path.join(UPLOAD_DIR, safe_name)
|
||||
|
||||
# Write in chunks to avoid doubling memory usage for large files
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(file_content)
|
||||
for offset in range(0, len(file_content), _UPLOAD_CHUNK_SIZE):
|
||||
f.write(file_content[offset : offset + _UPLOAD_CHUNK_SIZE])
|
||||
|
||||
# Count PDF pages
|
||||
page_count = _count_pages(file_content)
|
||||
|
|
@ -54,10 +63,12 @@ async def upload(filename: str, content_type: str, file_content: bytes) -> Docum
|
|||
|
||||
|
||||
async def find_all() -> list[Document]:
|
||||
"""Return all documents, newest first."""
|
||||
return await document_repo.find_all()
|
||||
|
||||
|
||||
async def find_by_id(doc_id: str) -> Document | None:
|
||||
"""Find a document by its ID, or return None."""
|
||||
return await document_repo.find_by_id(doc_id)
|
||||
|
||||
|
||||
|
|
@ -78,8 +89,12 @@ async def delete(doc_id: str) -> bool:
|
|||
os.unlink(real_path)
|
||||
elif os.path.exists(doc.storage_path):
|
||||
logger.warning("Refused to delete file outside upload dir: %s", doc.storage_path)
|
||||
except FileNotFoundError:
|
||||
logger.info("File already removed: %s", doc.storage_path)
|
||||
except PermissionError:
|
||||
logger.error("Permission denied deleting file: %s", doc.storage_path)
|
||||
except OSError:
|
||||
logger.warning("Could not delete file: %s", doc.storage_path)
|
||||
logger.warning("Could not delete file: %s", doc.storage_path, exc_info=True)
|
||||
|
||||
return await document_repo.delete(doc_id)
|
||||
|
||||
|
|
@ -100,6 +115,9 @@ def _count_pages(file_content: bytes) -> int | None:
|
|||
try:
|
||||
info = pdfinfo_from_bytes(file_content)
|
||||
return info.get("Pages")
|
||||
except Exception:
|
||||
logger.warning("Could not count pages", exc_info=True)
|
||||
except (FileNotFoundError, OSError) as exc:
|
||||
logger.warning("Could not count pages: %s", exc)
|
||||
return None
|
||||
except Exception:
|
||||
logger.warning("Unexpected error counting pages", exc_info=True)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -1,13 +1,13 @@
|
|||
"""Tests for AnalysisService — focus on _on_task_done callback (bug #1 fix)."""
|
||||
"""Tests for AnalysisService — callbacks, concurrency, and orchestration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from services.analysis_service import _on_task_done
|
||||
from services.analysis_service import AnalysisService, _on_task_done
|
||||
|
||||
|
||||
class TestOnTaskDone:
|
||||
|
|
@ -69,3 +69,46 @@ class TestOnTaskDone:
|
|||
await asyncio.sleep(0)
|
||||
|
||||
mock_mark.assert_not_called()
|
||||
|
||||
|
||||
class TestAnalysisServiceConcurrency:
|
||||
"""Verify that the semaphore limits concurrent analysis jobs."""
|
||||
|
||||
def test_semaphore_initialized_with_max_concurrent(self):
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter, max_concurrent=5)
|
||||
assert service._semaphore._value == 5
|
||||
|
||||
def test_default_max_concurrent(self):
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter)
|
||||
assert service._semaphore._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semaphore_limits_parallel_jobs(self):
|
||||
"""Only max_concurrent jobs should run in parallel; others must wait."""
|
||||
call_order: list[str] = []
|
||||
blocker = asyncio.Event()
|
||||
|
||||
converter = MagicMock()
|
||||
service = AnalysisService(converter=converter, max_concurrent=1)
|
||||
|
||||
async def fake_inner(self, *args, **kwargs):
|
||||
call_order.append("start")
|
||||
await blocker.wait()
|
||||
call_order.append("end")
|
||||
|
||||
with patch.object(AnalysisService, "_run_analysis_inner", fake_inner):
|
||||
t1 = asyncio.create_task(service._run_analysis("j1", "/f", "f.pdf"))
|
||||
t2 = asyncio.create_task(service._run_analysis("j2", "/f", "f.pdf"))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
# With max_concurrent=1, only one task should have started
|
||||
assert call_order.count("start") == 1
|
||||
|
||||
blocker.set()
|
||||
await asyncio.gather(t1, t2)
|
||||
|
||||
# Both should have completed
|
||||
assert call_order.count("start") == 2
|
||||
assert call_order.count("end") == 2
|
||||
|
|
|
|||
|
|
@ -29,8 +29,9 @@ class TestHealthEndpoint:
|
|||
resp = client.get("/api/health")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["status"] == "ok"
|
||||
assert data["status"] in ("ok", "degraded")
|
||||
assert "engine" in data
|
||||
assert "database" in data
|
||||
|
||||
|
||||
class TestDocumentEndpoints:
|
||||
|
|
@ -102,7 +103,20 @@ class TestDocumentEndpoints:
|
|||
"/api/documents/upload",
|
||||
files={"file": ("big.pdf", b"x", "application/pdf")},
|
||||
)
|
||||
assert resp.status_code == 413
|
||||
assert resp.status_code == 400
|
||||
|
||||
@patch("services.document_service.find_by_id", new_callable=AsyncMock)
|
||||
def test_preview_page_out_of_range(self, mock_find, client):
|
||||
mock_find.return_value = Document(
|
||||
id="d1",
|
||||
filename="test.pdf",
|
||||
page_count=3,
|
||||
storage_path="/tmp/test.pdf",
|
||||
)
|
||||
|
||||
resp = client.get("/api/documents/d1/preview?page=10")
|
||||
assert resp.status_code == 400
|
||||
assert "out of range" in resp.json()["detail"]
|
||||
|
||||
@patch("services.document_service.delete", new_callable=AsyncMock)
|
||||
def test_delete_document(self, mock_delete, client):
|
||||
|
|
|
|||
137
document-parser/tests/test_document_service.py
Normal file
137
document-parser/tests/test_document_service.py
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
"""Tests for document_service — upload, preview, page counting, and deletion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from domain.models import Document
|
||||
from services import document_service
|
||||
|
||||
|
||||
class TestUploadValidation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_oversized_file(self):
|
||||
content = b"x" * (document_service.MAX_FILE_SIZE + 1)
|
||||
with pytest.raises(ValueError, match="File too large"):
|
||||
await document_service.upload("big.pdf", "application/pdf", content)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_non_pdf(self):
|
||||
content = b"NOT-A-PDF-FILE"
|
||||
with pytest.raises(ValueError, match="not a PDF"):
|
||||
await document_service.upload("fake.pdf", "application/pdf", content)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_accepts_valid_pdf(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(document_service, "UPLOAD_DIR", str(tmp_path))
|
||||
|
||||
mock_insert = AsyncMock()
|
||||
with (
|
||||
patch("persistence.document_repo.insert", mock_insert),
|
||||
patch.object(document_service, "_count_pages", return_value=5),
|
||||
):
|
||||
content = b"%PDF-1.4 fake pdf content"
|
||||
doc = await document_service.upload("test.pdf", "application/pdf", content)
|
||||
|
||||
assert doc.filename == "test.pdf"
|
||||
assert doc.file_size == len(content)
|
||||
assert doc.page_count == 5
|
||||
mock_insert.assert_called_once()
|
||||
|
||||
# Verify file was actually written to disk
|
||||
assert os.path.exists(doc.storage_path)
|
||||
with open(doc.storage_path, "rb") as f:
|
||||
assert f.read() == content
|
||||
|
||||
|
||||
class TestGeneratePreview:
|
||||
def test_raises_on_invalid_page(self):
|
||||
"""generate_preview should raise ValueError when page is out of range."""
|
||||
with (
|
||||
patch("services.document_service.convert_from_bytes", return_value=[]),
|
||||
pytest.raises(ValueError, match="Page 1 not found"),
|
||||
):
|
||||
document_service.generate_preview(b"%PDF-fake", page=1)
|
||||
|
||||
def test_returns_png_bytes(self):
|
||||
"""generate_preview should return PNG bytes from pdf2image."""
|
||||
mock_image = MagicMock()
|
||||
mock_image.save = MagicMock(side_effect=lambda buf, format: buf.write(b"PNG-DATA"))
|
||||
|
||||
with patch("services.document_service.convert_from_bytes", return_value=[mock_image]):
|
||||
result = document_service.generate_preview(b"%PDF-fake", page=1, dpi=72)
|
||||
|
||||
assert result == b"PNG-DATA"
|
||||
|
||||
|
||||
class TestCountPages:
|
||||
def test_returns_page_count(self):
|
||||
with patch(
|
||||
"services.document_service.pdfinfo_from_bytes",
|
||||
return_value={"Pages": 42},
|
||||
):
|
||||
assert document_service._count_pages(b"pdf") == 42
|
||||
|
||||
def test_returns_none_on_error(self):
|
||||
with patch(
|
||||
"services.document_service.pdfinfo_from_bytes",
|
||||
side_effect=FileNotFoundError("poppler not found"),
|
||||
):
|
||||
assert document_service._count_pages(b"pdf") is None
|
||||
|
||||
|
||||
class TestDelete:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_removes_file_and_records(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(document_service, "UPLOAD_DIR", str(tmp_path))
|
||||
|
||||
# Create a fake file
|
||||
fake_file = tmp_path / "test.pdf"
|
||||
fake_file.write_bytes(b"content")
|
||||
|
||||
doc = Document(
|
||||
id="doc-1",
|
||||
filename="test.pdf",
|
||||
storage_path=str(fake_file),
|
||||
)
|
||||
|
||||
with (
|
||||
patch("persistence.document_repo.find_by_id", AsyncMock(return_value=doc)),
|
||||
patch("persistence.analysis_repo.delete_by_document", AsyncMock(return_value=2)),
|
||||
patch("persistence.document_repo.delete", AsyncMock(return_value=True)),
|
||||
):
|
||||
result = await document_service.delete("doc-1")
|
||||
|
||||
assert result is True
|
||||
assert not fake_file.exists()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_refuses_file_outside_upload_dir(self, tmp_path, monkeypatch):
|
||||
"""Files outside UPLOAD_DIR should not be deleted (path traversal protection)."""
|
||||
monkeypatch.setattr(document_service, "UPLOAD_DIR", str(tmp_path / "uploads"))
|
||||
os.makedirs(tmp_path / "uploads", exist_ok=True)
|
||||
|
||||
# File is outside the upload dir
|
||||
outside_file = tmp_path / "secret.txt"
|
||||
outside_file.write_bytes(b"secret")
|
||||
|
||||
doc = Document(id="doc-1", filename="x.pdf", storage_path=str(outside_file))
|
||||
|
||||
with (
|
||||
patch("persistence.document_repo.find_by_id", AsyncMock(return_value=doc)),
|
||||
patch("persistence.analysis_repo.delete_by_document", AsyncMock(return_value=0)),
|
||||
patch("persistence.document_repo.delete", AsyncMock(return_value=True)),
|
||||
):
|
||||
await document_service.delete("doc-1")
|
||||
|
||||
# File should NOT have been deleted
|
||||
assert outside_file.exists()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_not_found_returns_false(self):
|
||||
with patch("persistence.document_repo.find_by_id", AsyncMock(return_value=None)):
|
||||
result = await document_service.delete("missing")
|
||||
assert result is False
|
||||
93
document-parser/tests/test_rate_limiter.py
Normal file
93
document-parser/tests/test_rate_limiter.py
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
"""Tests for the in-memory rate limiter middleware."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from infra.rate_limiter import RateLimiterMiddleware, _ClientBucket
|
||||
|
||||
|
||||
class TestClientBucket:
|
||||
def test_count_recent_filters_old_entries(self):
|
||||
bucket = _ClientBucket(timestamps=[1.0, 2.0, 3.0, 10.0])
|
||||
count = bucket.count_recent(window=5.0, now=12.0)
|
||||
assert count == 1 # only 10.0 is within [7.0, 12.0]
|
||||
|
||||
def test_count_recent_keeps_all_when_within_window(self):
|
||||
bucket = _ClientBucket(timestamps=[10.0, 11.0, 12.0])
|
||||
count = bucket.count_recent(window=60.0, now=15.0)
|
||||
assert count == 3
|
||||
|
||||
def test_add(self):
|
||||
bucket = _ClientBucket()
|
||||
bucket.add(1.0)
|
||||
bucket.add(2.0)
|
||||
assert len(bucket.timestamps) == 2
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def limited_app():
|
||||
"""FastAPI app with a very low rate limit for testing."""
|
||||
app = FastAPI()
|
||||
app.add_middleware(
|
||||
RateLimiterMiddleware,
|
||||
requests_per_window=3,
|
||||
window_seconds=60,
|
||||
exclude_paths=("/health",),
|
||||
)
|
||||
|
||||
@app.get("/test")
|
||||
def test_endpoint():
|
||||
return {"ok": True}
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(limited_app):
|
||||
return TestClient(limited_app)
|
||||
|
||||
|
||||
class TestRateLimiterMiddleware:
|
||||
def test_allows_requests_under_limit(self, client):
|
||||
for _ in range(3):
|
||||
resp = client.get("/test")
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_blocks_requests_over_limit(self, client):
|
||||
for _ in range(3):
|
||||
client.get("/test")
|
||||
|
||||
resp = client.get("/test")
|
||||
assert resp.status_code == 429
|
||||
assert resp.json()["detail"] == "Too many requests"
|
||||
assert "Retry-After" in resp.headers
|
||||
|
||||
def test_health_excluded_from_limit(self, client):
|
||||
# Exhaust the limit
|
||||
for _ in range(3):
|
||||
client.get("/test")
|
||||
|
||||
# Health should still work
|
||||
resp = client.get("/health")
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_window_resets(self, client):
|
||||
"""After the window expires, requests should be allowed again."""
|
||||
for _ in range(3):
|
||||
client.get("/test")
|
||||
|
||||
assert client.get("/test").status_code == 429
|
||||
|
||||
# Simulate time passing beyond the window
|
||||
with patch("time.monotonic", return_value=1e12):
|
||||
resp = client.get("/test")
|
||||
assert resp.status_code == 200
|
||||
75
document-parser/tests/test_settings.py
Normal file
75
document-parser/tests/test_settings.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
"""Tests for Settings — environment variable parsing and defaults."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from infra.settings import Settings
|
||||
|
||||
|
||||
class TestSettingsDefaults:
|
||||
def test_default_values(self):
|
||||
s = Settings()
|
||||
assert s.app_version == "dev"
|
||||
assert s.conversion_engine == "local"
|
||||
assert s.docling_serve_url == "http://localhost:5001"
|
||||
assert s.docling_serve_api_key is None
|
||||
assert s.conversion_timeout == 600
|
||||
assert s.upload_dir == "./uploads"
|
||||
assert s.db_path == "./data/docling_studio.db"
|
||||
assert "http://localhost:3000" in s.cors_origins
|
||||
|
||||
def test_frozen(self):
|
||||
"""Settings should be immutable."""
|
||||
import pytest
|
||||
|
||||
s = Settings()
|
||||
with pytest.raises(AttributeError):
|
||||
s.upload_dir = "/other" # type: ignore[misc]
|
||||
|
||||
|
||||
class TestSettingsFromEnv:
|
||||
def test_reads_env_vars(self, monkeypatch):
|
||||
monkeypatch.setenv("APP_VERSION", "1.2.3")
|
||||
monkeypatch.setenv("CONVERSION_ENGINE", "remote")
|
||||
monkeypatch.setenv("DOCLING_SERVE_URL", "http://serve:9000")
|
||||
monkeypatch.setenv("DOCLING_SERVE_API_KEY", "secret-key")
|
||||
monkeypatch.setenv("CONVERSION_TIMEOUT", "120")
|
||||
monkeypatch.setenv("UPLOAD_DIR", "/data/uploads")
|
||||
monkeypatch.setenv("DB_PATH", "/data/test.db")
|
||||
monkeypatch.setenv("CORS_ORIGINS", "http://a.com, http://b.com")
|
||||
|
||||
s = Settings.from_env()
|
||||
|
||||
assert s.app_version == "1.2.3"
|
||||
assert s.conversion_engine == "remote"
|
||||
assert s.docling_serve_url == "http://serve:9000"
|
||||
assert s.docling_serve_api_key == "secret-key"
|
||||
assert s.conversion_timeout == 120
|
||||
assert s.upload_dir == "/data/uploads"
|
||||
assert s.db_path == "/data/test.db"
|
||||
assert s.cors_origins == ["http://a.com", "http://b.com"]
|
||||
|
||||
def test_defaults_when_env_empty(self, monkeypatch):
|
||||
"""When no env vars set, from_env returns sensible defaults."""
|
||||
for key in (
|
||||
"APP_VERSION",
|
||||
"CONVERSION_ENGINE",
|
||||
"DOCLING_SERVE_URL",
|
||||
"DOCLING_SERVE_API_KEY",
|
||||
"CONVERSION_TIMEOUT",
|
||||
"UPLOAD_DIR",
|
||||
"DB_PATH",
|
||||
"CORS_ORIGINS",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
s = Settings.from_env()
|
||||
|
||||
assert s.app_version == "dev"
|
||||
assert s.conversion_engine == "local"
|
||||
assert s.conversion_timeout == 600
|
||||
|
||||
def test_cors_origins_split(self, monkeypatch):
|
||||
monkeypatch.setenv("CORS_ORIGINS", "http://a.com,http://b.com,http://c.com")
|
||||
s = Settings.from_env()
|
||||
assert len(s.cors_origins) == 3
|
||||
assert s.cors_origins[2] == "http://c.com"
|
||||
Loading…
Reference in a new issue