ontocast.tool.chunk.section_llm

LLM section-label backfill for chunk preparation.

BatchSectionClassification

Bases: BasePydanticModel

LLM output assigning a section label to each numbered excerpt.

Source code in ontocast/tool/chunk/section_llm.py

class BatchSectionClassification(BasePydanticModel):
    """LLM output assigning a section label to each numbered excerpt."""

    assignments: list[SectionLabelAssignment] = Field(default_factory=list)

ChunkSectionClassification

Bases: BasePydanticModel

LLM output mapping one excerpt to a canonical section label.

Source code in ontocast/tool/chunk/section_llm.py

class ChunkSectionClassification(BasePydanticModel):
    """LLM output mapping one excerpt to a canonical section label."""

    label: str | None = Field(
        default=None,
        description="Canonical section label or null if not classifiable",
    )

SectionLabelAssignment

Bases: BasePydanticModel

One excerpt index and the section label assigned to it.

Source code in ontocast/tool/chunk/section_llm.py

class SectionLabelAssignment(BasePydanticModel):
    """One excerpt index and the section label assigned to it."""

    index: int = Field(description="Index of the excerpt, as given in the prompt")
    label: str | None = Field(
        default=None,
        description="Canonical section label or null if not classifiable",
    )

classify_section_with_llm(text, tools, schema, *, document_type_hint=None) async

Classify a text fragment with the section-label LLM prompt.

Source code in ontocast/tool/chunk/section_llm.py

async def classify_section_with_llm(
    text: str,
    tools: "ToolBox",
    schema: SectionLabelSchema,
    *,
    document_type_hint: str | None = None,
) -> str | None:
    """Classify a text fragment with the section-label LLM prompt."""
    fragment = fragment_for_text(text)
    if not fragment:
        return None
    parser = PydanticOutputParser(pydantic_object=ChunkSectionClassification)
    allowed = ", ".join(canonical_labels(schema))
    prompt = CHUNK_SECTION_CLASSIFICATION_PROMPT.format_prompt(
        allowed_labels=allowed,
        format_instructions=parser.get_format_instructions(),
        document_context=document_type_context(document_type_hint),
        fragment=fragment,
    )
    response = await tools.llm(prompt)
    parsed = parser.parse(response.content or "")
    return normalise_llm_label(parsed.label, schema)

classify_sections_batched(items, tools, schema, *, document_type_hint=None, batch_size=40) async

Classify many excerpts in as few LLM calls as possible.

One call covers up to batch_size excerpts, versus one call per excerpt for :func:classify_section_with_llm. Passing the excerpts together also gives the model the document's shape, which a single fragment cannot show.

Parameters:

Name	Type	Description	Default
items	list[tuple[int, str]]	(index, fragment) pairs in document order.	required
tools	'ToolBox'	ToolBox providing the LLM.	required
schema	SectionLabelSchema	Active section label schema.	required
document_type_hint	str | None	Optional free-text document type.	None
batch_size	int	Maximum excerpts per LLM call.	40

Returns:

Type	Description
dict[int, str | None] | None	Mapping of index to label (None where unclassifiable), or None
dict[int, str | None] | None	when the model's response could not be used, so the caller can fall
dict[int, str | None] | None	back to per-excerpt classification.

Source code in ontocast/tool/chunk/section_llm.py

async def classify_sections_batched(
    items: list[tuple[int, str]],
    tools: "ToolBox",
    schema: SectionLabelSchema,
    *,
    document_type_hint: str | None = None,
    batch_size: int = 40,
) -> dict[int, str | None] | None:
    """Classify many excerpts in as few LLM calls as possible.

    One call covers up to ``batch_size`` excerpts, versus one call per excerpt
    for :func:`classify_section_with_llm`. Passing the excerpts together also
    gives the model the document's shape, which a single fragment cannot show.

    Args:
        items: ``(index, fragment)`` pairs in document order.
        tools: ToolBox providing the LLM.
        schema: Active section label schema.
        document_type_hint: Optional free-text document type.
        batch_size: Maximum excerpts per LLM call.

    Returns:
        Mapping of index to label (``None`` where unclassifiable), or ``None``
        when the model's response could not be used, so the caller can fall
        back to per-excerpt classification.
    """
    if not items:
        return {}
    parser = PydanticOutputParser(pydantic_object=BatchSectionClassification)
    allowed = ", ".join(canonical_labels(schema))
    resolved: dict[int, str | None] = {}
    size = max(1, batch_size)

    for start in range(0, len(items), size):
        batch = items[start : start + size]
        prompt = CHUNK_SECTION_BATCH_CLASSIFICATION_PROMPT.format_prompt(
            allowed_labels=allowed,
            format_instructions=parser.get_format_instructions(),
            document_context=document_type_context(document_type_hint),
            items=format_batch_items(batch),
        )
        try:
            response = await tools.llm(prompt)
            parsed = parser.parse(response.content or "")
        except Exception as exc:
            logger.warning(
                "Batched section classification failed for %s excerpt(s): %s",
                len(batch),
                exc,
            )
            return None
        known = {index for index, _ in batch}
        for assignment in parsed.assignments:
            if assignment.index in known:
                resolved[assignment.index] = normalise_llm_label(
                    assignment.label, schema
                )
    return resolved

fragment_for_text(text)

Return a short excerpt suitable for LLM section classification.

Source code in ontocast/tool/chunk/section_llm.py

def fragment_for_text(text: str) -> str:
    """Return a short excerpt suitable for LLM section classification."""
    for line in text.splitlines():
        stripped = line.strip()
        if stripped.startswith("#"):
            return stripped[:_FRAGMENT_MAX_CHARS]
    snippet = text.strip()
    return snippet[:_FRAGMENT_MAX_CHARS]

llm_backfill_section_labels(segments, tools, *, section_schema_id=None, document_type_hint=None, section_tag_min_chars=80, batch_size=40, schema=None) async

Set section_label on segments that are still unlabeled (mutates in place).

Classifies in batches when batch_size is positive, falling back to one call per segment if the batched response cannot be used.

Parameters:

Name	Type	Description	Default
segments	list	Prepare segments, mutated in place.	required
tools	'ToolBox'	ToolBox providing the LLM.	required
section_schema_id	str | None	Raw request value; used only when schema is not given.	None
document_type_hint	str | None	Free-text document type, also passed to the prompt.	None
section_tag_min_chars	int	Minimum segment length to be worth classifying.	80
batch_size	int	Excerpts per LLM call; 0 restores one call per segment.	40
schema	SectionLabelSchema | None	Already-resolved schema. Callers that resolved it themselves must pass it: re-deriving from the raw request would ignore a text-based schema detection, and labels outside the re-derived schema are silently discarded.	None

Source code in ontocast/tool/chunk/section_llm.py

async def llm_backfill_section_labels(
    segments: list,
    tools: "ToolBox",
    *,
    section_schema_id: str | None = None,
    document_type_hint: str | None = None,
    section_tag_min_chars: int = 80,
    batch_size: int = 40,
    schema: SectionLabelSchema | None = None,
) -> None:
    """Set ``section_label`` on segments that are still unlabeled (mutates in place).

    Classifies in batches when ``batch_size`` is positive, falling back to one
    call per segment if the batched response cannot be used.

    Args:
        segments: Prepare segments, mutated in place.
        tools: ToolBox providing the LLM.
        section_schema_id: Raw request value; used only when ``schema`` is not
            given.
        document_type_hint: Free-text document type, also passed to the prompt.
        section_tag_min_chars: Minimum segment length to be worth classifying.
        batch_size: Excerpts per LLM call; 0 restores one call per segment.
        schema: Already-resolved schema. Callers that resolved it themselves
            **must** pass it: re-deriving from the raw request would ignore a
            text-based schema detection, and labels outside the re-derived
            schema are silently discarded.
    """
    if schema is None:
        schema = load_section_label_schema(
            resolve_section_schema_id(
                section_schema_id=section_schema_id,
                document_type_hint=document_type_hint,
            )
        )
    min_chars = max(0, section_tag_min_chars)

    def _needs_llm_backfill(index: int) -> bool:
        segment = segments[index]
        if segment.section_label is not None:
            return False
        text = segment.text.strip()
        fragment = fragment_for_text(segment.text)
        if not fragment:
            return False
        if len(text) >= min_chars:
            return True
        if fragment.lstrip().startswith("#"):
            return True
        return bool(segment.headings)

    unlabeled_indices = [
        index for index in range(len(segments)) if _needs_llm_backfill(index)
    ]
    if not unlabeled_indices:
        return

    if batch_size > 0:
        batched = await classify_sections_batched(
            [
                (
                    index,
                    fragment_for_text(segments[index].text)[:_BATCH_FRAGMENT_MAX_CHARS],
                )
                for index in unlabeled_indices
            ],
            tools,
            schema,
            document_type_hint=document_type_hint,
            batch_size=batch_size,
        )
        if batched is not None:
            _apply_llm_labels(segments, batched)
            return
        logger.info(
            "Falling back to per-segment section classification for %s segment(s)",
            len(unlabeled_indices),
        )

    worker_limit = max(1, tools.config.server.parallel_workers)
    semaphore = asyncio.Semaphore(worker_limit)

    async def classify_index(index: int) -> tuple[int, str | None]:
        wait_start = time.perf_counter()
        async with semaphore:
            record_active_span(
                "chunk section classify/worker_wait", time.perf_counter() - wait_start
            )
            segment = segments[index]
            try:
                label = await classify_section_with_llm(
                    segment.text,
                    tools,
                    schema,
                    document_type_hint=document_type_hint,
                )
                return index, label
            except Exception as exc:
                logger.warning(
                    "LLM section classification failed for segment %s: %s",
                    index,
                    exc,
                )
                return index, None

    # classify_index catches its own errors, but the gather must not abort the
    # whole backfill if one slips through -- an unlabeled segment is survivable,
    # an unchunked document is not.
    results = await asyncio.gather(
        *[classify_index(index) for index in unlabeled_indices],
        return_exceptions=True,
    )
    labels = {
        index: label
        for index, label in (
            item for item in results if not isinstance(item, BaseException)
        )
    }
    for item in results:
        if isinstance(item, BaseException):
            logger.warning("Section classification task failed: %s", item)
    _apply_llm_labels(segments, labels)