rag/types.py

"""Shapes of the files the ingestion writes."""

from typing import TypedDict


class Page(TypedDict):
    """One chapter, as listed in a book's SUMMARY.md."""
    book: str
    path: str  # relative to the book's source directory, e.g. "ch10-01-syntax.md"
    title: str
    part: str | None  # the mdBook part heading above it
    url: str  # published HTML, without a fragment


class ManifestBook(TypedDict):
    title: str
    repo: str
    sha: str
    base_url: str
    license: str
    excluded_parts: list[str]
    page_count: int
    pages: list[Page]


class ChunkMetadata(TypedDict):
    book: str
    book_title: str
    part: str
    chapter: str
    heading_path: list[str]
    path: str
    url: str  # published HTML, with the section fragment when there is one
    has_code: bool
    code_tags: list[str]


class Chunk(TypedDict):
    id: str
    text: str
    metadata: ChunkMetadata


class EvalRow(TypedDict):
    """One question/chunk pair in `eval/dataset.json`.

    `chunk_id` is the chunk the question was generated from, and what the retriever
    is scored against.
    """

    question: str
    chunk_id: str
    book: str