Add chunking service orchestration, API endpoints, and wiring
AnalysisService gains rechunk() and inline chunking during conversion. ChunkingOptionsRequest/ChunkResponse schemas, POST rechunk endpoint, and conditional chunker injection in main.py (local engine only).
This commit is contained in:
parent
4b1ec364f4
commit
a9517d38eb
4 changed files with 153 additions and 10 deletions
|
|
@ -7,7 +7,7 @@ from typing import Annotated
|
|||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
|
||||
from api.schemas import AnalysisResponse, CreateAnalysisRequest
|
||||
from api.schemas import AnalysisResponse, ChunkResponse, CreateAnalysisRequest, RechunkRequest
|
||||
from services.analysis_service import AnalysisService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -30,6 +30,8 @@ def _to_response(job) -> AnalysisResponse:
|
|||
content_markdown=job.content_markdown,
|
||||
content_html=job.content_html,
|
||||
pages_json=job.pages_json,
|
||||
chunks_json=job.chunks_json,
|
||||
has_document_json=job.document_json is not None,
|
||||
error_message=job.error_message,
|
||||
started_at=str(job.started_at) if job.started_at else None,
|
||||
completed_at=str(job.completed_at) if job.completed_at else None,
|
||||
|
|
@ -47,8 +49,16 @@ async def create_analysis(body: CreateAnalysisRequest, service: ServiceDep):
|
|||
if body.pipelineOptions:
|
||||
pipeline_opts = body.pipelineOptions.model_dump()
|
||||
|
||||
chunking_opts = None
|
||||
if body.chunkingOptions:
|
||||
chunking_opts = body.chunkingOptions.model_dump()
|
||||
|
||||
try:
|
||||
job = await service.create(body.documentId, pipeline_options=pipeline_opts)
|
||||
job = await service.create(
|
||||
body.documentId,
|
||||
pipeline_options=pipeline_opts,
|
||||
chunking_options=chunking_opts,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e)) from e
|
||||
|
||||
|
|
@ -71,6 +81,24 @@ async def get_analysis(job_id: str, service: ServiceDep):
|
|||
return _to_response(job)
|
||||
|
||||
|
||||
@router.post("/{job_id}/rechunk", response_model=list[ChunkResponse])
|
||||
async def rechunk_analysis(job_id: str, body: RechunkRequest, service: ServiceDep):
|
||||
"""Re-chunk a completed analysis with new chunking options."""
|
||||
try:
|
||||
chunks = await service.rechunk(job_id, body.chunkingOptions.model_dump())
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from e
|
||||
return [
|
||||
ChunkResponse(
|
||||
text=c.text,
|
||||
headings=c.headings,
|
||||
source_page=c.source_page,
|
||||
token_count=c.token_count,
|
||||
)
|
||||
for c in chunks
|
||||
]
|
||||
|
||||
|
||||
@router.delete("/{job_id}", status_code=204)
|
||||
async def delete_analysis(job_id: str, service: ServiceDep):
|
||||
"""Delete an analysis job."""
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ def _to_camel(name: str) -> str:
|
|||
|
||||
class _CamelModel(BaseModel):
|
||||
"""Base model that serializes field names to camelCase."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
alias_generator=_to_camel,
|
||||
populate_by_name=True,
|
||||
|
|
@ -42,6 +43,8 @@ class AnalysisResponse(_CamelModel):
|
|||
content_markdown: str | None = None
|
||||
content_html: str | None = None
|
||||
pages_json: str | None = None
|
||||
chunks_json: str | None = None
|
||||
has_document_json: bool = False
|
||||
error_message: str | None = None
|
||||
started_at: str | datetime | None = None
|
||||
completed_at: str | datetime | None = None
|
||||
|
|
@ -50,6 +53,7 @@ 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"
|
||||
|
|
@ -76,6 +80,41 @@ class PipelineOptionsRequest(BaseModel):
|
|||
return v
|
||||
|
||||
|
||||
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
|
||||
|
||||
@field_validator("chunker_type")
|
||||
@classmethod
|
||||
def validate_chunker_type(cls, v: str) -> str:
|
||||
if v not in ("hybrid", "hierarchical"):
|
||||
raise ValueError('chunker_type must be "hybrid" or "hierarchical"')
|
||||
return v
|
||||
|
||||
@field_validator("max_tokens")
|
||||
@classmethod
|
||||
def validate_max_tokens(cls, v: int) -> int:
|
||||
if v < 64 or v > 8192:
|
||||
raise ValueError("max_tokens must be between 64 and 8192")
|
||||
return v
|
||||
|
||||
|
||||
class ChunkResponse(_CamelModel):
|
||||
text: str
|
||||
headings: list[str] = []
|
||||
source_page: int | None = None
|
||||
token_count: int = 0
|
||||
|
||||
|
||||
class CreateAnalysisRequest(BaseModel):
|
||||
documentId: str # camelCase to match existing frontend contract
|
||||
pipelineOptions: PipelineOptionsRequest | None = None
|
||||
chunkingOptions: ChunkingOptionsRequest | None = None
|
||||
|
||||
|
||||
class RechunkRequest(BaseModel):
|
||||
chunkingOptions: ChunkingOptionsRequest
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ def _build_converter():
|
|||
"""Build the converter adapter based on configuration."""
|
||||
if settings.conversion_engine == "remote":
|
||||
from infra.serve_converter import ServeConverter
|
||||
|
||||
logger.info("Using remote Docling Serve at %s", settings.docling_serve_url)
|
||||
return ServeConverter(
|
||||
base_url=settings.docling_serve_url,
|
||||
|
|
@ -47,14 +48,26 @@ def _build_converter():
|
|||
)
|
||||
else:
|
||||
from infra.local_converter import LocalConverter
|
||||
|
||||
logger.info("Using local Docling converter")
|
||||
return LocalConverter()
|
||||
|
||||
|
||||
def _build_chunker():
|
||||
"""Build the chunker adapter — only available in local mode."""
|
||||
if settings.conversion_engine == "local":
|
||||
from infra.local_chunker import LocalChunker
|
||||
|
||||
return LocalChunker()
|
||||
return None
|
||||
|
||||
|
||||
def _build_analysis_service() -> AnalysisService:
|
||||
converter = _build_converter()
|
||||
chunker = _build_chunker()
|
||||
return AnalysisService(
|
||||
converter=converter,
|
||||
chunker=chunker,
|
||||
conversion_timeout=settings.conversion_timeout,
|
||||
)
|
||||
|
||||
|
|
@ -63,6 +76,7 @@ def _build_analysis_service() -> AnalysisService:
|
|||
# FastAPI app
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
await init_db()
|
||||
|
|
|
|||
|
|
@ -13,24 +13,46 @@ import logging
|
|||
from dataclasses import asdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from domain.models import AnalysisJob
|
||||
from domain.value_objects import ConversionOptions, ConversionResult
|
||||
from domain.models import AnalysisJob, AnalysisStatus
|
||||
from domain.value_objects import ChunkingOptions, ChunkResult, ConversionOptions, ConversionResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from domain.ports import DocumentConverter
|
||||
from domain.ports import DocumentChunker, DocumentConverter
|
||||
from persistence import analysis_repo, document_repo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _chunk_to_dict(c: ChunkResult) -> dict:
|
||||
"""Serialize ChunkResult to a camelCase dict matching the frontend API contract."""
|
||||
return {
|
||||
"text": c.text,
|
||||
"headings": c.headings,
|
||||
"sourcePage": c.source_page,
|
||||
"tokenCount": c.token_count,
|
||||
}
|
||||
|
||||
|
||||
class AnalysisService:
|
||||
"""Orchestrates document analysis using an injected converter."""
|
||||
|
||||
def __init__(self, converter: DocumentConverter, conversion_timeout: int = 600):
|
||||
def __init__(
|
||||
self,
|
||||
converter: DocumentConverter,
|
||||
chunker: DocumentChunker | None = None,
|
||||
conversion_timeout: int = 600,
|
||||
):
|
||||
self._converter = converter
|
||||
self._chunker = chunker
|
||||
self._conversion_timeout = conversion_timeout
|
||||
|
||||
async def create(self, document_id: str, *, pipeline_options: dict | None = None) -> AnalysisJob:
|
||||
async def create(
|
||||
self,
|
||||
document_id: str,
|
||||
*,
|
||||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> AnalysisJob:
|
||||
"""Create a new analysis job and launch background processing."""
|
||||
doc = await document_repo.find_by_id(document_id)
|
||||
if not doc:
|
||||
|
|
@ -41,7 +63,13 @@ class AnalysisService:
|
|||
await analysis_repo.insert(job)
|
||||
|
||||
task = asyncio.create_task(
|
||||
self._run_analysis(job.id, doc.storage_path, doc.filename, pipeline_options)
|
||||
self._run_analysis(
|
||||
job.id,
|
||||
doc.storage_path,
|
||||
doc.filename,
|
||||
pipeline_options,
|
||||
chunking_options,
|
||||
)
|
||||
)
|
||||
task.add_done_callback(functools.partial(_on_task_done, job_id=job.id))
|
||||
|
||||
|
|
@ -56,10 +84,35 @@ class AnalysisService:
|
|||
async def delete(self, job_id: str) -> bool:
|
||||
return await analysis_repo.delete(job_id)
|
||||
|
||||
async def rechunk(self, job_id: str, chunking_options: dict) -> list[ChunkResult]:
|
||||
"""Re-chunk an existing completed analysis with new options."""
|
||||
job = await analysis_repo.find_by_id(job_id)
|
||||
if not job:
|
||||
raise ValueError(f"Analysis not found: {job_id}")
|
||||
if job.status != AnalysisStatus.COMPLETED:
|
||||
raise ValueError(f"Analysis is not completed: {job_id}")
|
||||
if not job.document_json:
|
||||
raise ValueError(f"No document data available for re-chunking: {job_id}")
|
||||
if not self._chunker:
|
||||
raise ValueError("Chunking is not available")
|
||||
|
||||
options = ChunkingOptions(**chunking_options)
|
||||
chunks = await self._chunker.chunk(job.document_json, options)
|
||||
|
||||
chunks_json = json.dumps([_chunk_to_dict(c) for c in chunks])
|
||||
await analysis_repo.update_chunks(job_id, chunks_json)
|
||||
|
||||
return chunks
|
||||
|
||||
async def _run_analysis(
|
||||
self, job_id: str, file_path: str, filename: str, pipeline_options: dict | None = None,
|
||||
self,
|
||||
job_id: str,
|
||||
file_path: str,
|
||||
filename: str,
|
||||
pipeline_options: dict | None = None,
|
||||
chunking_options: dict | None = None,
|
||||
) -> None:
|
||||
"""Background task: run conversion and update job status."""
|
||||
"""Background task: run conversion and optionally chunk."""
|
||||
try:
|
||||
job = await analysis_repo.find_by_id(job_id)
|
||||
if not job:
|
||||
|
|
@ -79,10 +132,19 @@ class AnalysisService:
|
|||
|
||||
pages_json = json.dumps([asdict(p) for p in result.pages])
|
||||
|
||||
chunks_json = None
|
||||
if chunking_options and self._chunker and result.document_json:
|
||||
chunk_opts = ChunkingOptions(**chunking_options)
|
||||
chunks = await self._chunker.chunk(result.document_json, chunk_opts)
|
||||
chunks_json = json.dumps([_chunk_to_dict(c) for c in chunks])
|
||||
logger.info("Chunking produced %d chunks for job %s", len(chunks), job_id)
|
||||
|
||||
job.mark_completed(
|
||||
markdown=result.content_markdown,
|
||||
html=result.content_html,
|
||||
pages_json=pages_json,
|
||||
document_json=result.document_json,
|
||||
chunks_json=chunks_json,
|
||||
)
|
||||
await analysis_repo.update_status(job)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue