diff --git a/src/generation/answerer.py b/src/generation/answerer.py index a4ea4fa..b2a43be 100644 --- a/src/generation/answerer.py +++ b/src/generation/answerer.py @@ -45,6 +45,9 @@ def build_prompt(question: str, results: list[SearchResult]) -> list[dict]: def _source_name(metadata: dict[str, str | int]) -> str: """Build a display name for a source from its metadata.""" path = metadata.get("relative_path", metadata.get("filename", "unknown")) + section = metadata.get("section_header") + if section: + return f"Source: {path} | Section: {section}" page = metadata.get("page_number", "?") return f"Source: {path} | Page {page}" diff --git a/src/ingestion/chunker.py b/src/ingestion/chunker.py index 995f2c8..394d82d 100644 --- a/src/ingestion/chunker.py +++ b/src/ingestion/chunker.py @@ -1,12 +1,23 @@ import hashlib -from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_text_splitters import ( + MarkdownHeaderTextSplitter, + RecursiveCharacterTextSplitter, +) from src.models import Chunk, Document -# --- LangChain boundary: RecursiveCharacterTextSplitter used for text splitting. -# All other pipeline stages (embeddings, vector store, LLM calls) use -# direct library calls, not LangChain. --- +# --- LangChain boundary: RecursiveCharacterTextSplitter and +# MarkdownHeaderTextSplitter used for text splitting. +# All other pipeline stages (embeddings, vector store, LLM calls) +# use direct library calls, not LangChain. --- + +MARKDOWN_HEADERS = [ + ("#", "h1"), + ("##", "h2"), + ("###", "h3"), + ("####", "h4"), +] def chunk_documents( @@ -16,24 +27,23 @@ def chunk_documents( chunk_overlap: int = 100, ) -> list[Chunk]: """Split documents into chunks with preserved and extended metadata.""" - splitter = RecursiveCharacterTextSplitter( - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - length_function=len, - ) - chunks: list[Chunk] = [] for doc in documents: doc_hash = hashlib.md5(doc.text.encode()).hexdigest() - splits = splitter.split_text(doc.text) - for i, text in enumerate(splits): + if doc.metadata.get("doc_type") == "md": + doc_chunks = _chunk_markdown(doc, chunk_size, chunk_overlap) + else: + doc_chunks = _chunk_text(doc, chunk_size, chunk_overlap) + + for i, (text, extra_meta) in enumerate(doc_chunks): chunks.append( Chunk( text=text, metadata={ **doc.metadata, + **extra_meta, "chunk_index": i, "doc_hash": doc_hash, }, @@ -41,3 +51,52 @@ def chunk_documents( ) return chunks + + +def _chunk_text( + doc: Document, chunk_size: int, chunk_overlap: int +) -> list[tuple[str, dict]]: + """Split non-Markdown text using RecursiveCharacterTextSplitter.""" + splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + length_function=len, + ) + splits = splitter.split_text(doc.text) + return [(text, {}) for text in splits] + + +def _chunk_markdown( + doc: Document, chunk_size: int, chunk_overlap: int +) -> list[tuple[str, dict]]: + """Split Markdown by headers first, then by size.""" + md_splitter = MarkdownHeaderTextSplitter( + headers_to_split_on=MARKDOWN_HEADERS, + strip_headers=False, + ) + header_splits = md_splitter.split_text(doc.text) + + size_splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + length_function=len, + ) + + result: list[tuple[str, dict]] = [] + for split in header_splits: + section_header = _build_section_header(split.metadata) + sub_splits = size_splitter.split_text(split.page_content) + for text in sub_splits: + extra = {"section_header": section_header} if section_header else {} + result.append((text, extra)) + + return result + + +def _build_section_header(metadata: dict[str, str]) -> str: + """Build a section header string like 'Guide > Setup > Prerequisites'.""" + parts = [] + for key in ("h1", "h2", "h3", "h4"): + if key in metadata: + parts.append(metadata[key]) + return " > ".join(parts) diff --git a/src/ingestion/loader.py b/src/ingestion/loader.py index dae6742..2294e8e 100644 --- a/src/ingestion/loader.py +++ b/src/ingestion/loader.py @@ -1,41 +1,88 @@ from pathlib import Path -from langchain_community.document_loaders import PyPDFLoader +from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader from src.models import Document -# --- LangChain boundary: PyPDFLoader used for PDF parsing. -# This is one of two LangChain touch-points (the other is -# RecursiveCharacterTextSplitter in chunker.py). All other -# pipeline stages use direct library calls. --- +# --- LangChain boundary: PyPDFLoader, Docx2txtLoader, and TextLoader used +# for document parsing. This is one of two LangChain touch-points +# (the other is text splitters in chunker.py). All other pipeline +# stages use direct library calls. --- + +SUPPORTED_EXTENSIONS = ("*.pdf", "*.docx", "*.txt", "*.md") + +EXTENSION_TO_DOC_TYPE = { + ".pdf": "pdf", + ".docx": "docx", + ".txt": "txt", + ".md": "md", +} def load_folder(path: Path, *, recursive: bool = True) -> list[Document]: - """Load all PDF files from a folder into Document objects.""" + """Load all supported documents from a folder.""" path = Path(path) - if recursive: - pdf_files = sorted(path.rglob("*.pdf")) - else: - pdf_files = sorted(path.glob("*.pdf")) + files: list[Path] = [] + for ext in SUPPORTED_EXTENSIONS: + if recursive: + files.extend(path.rglob(ext)) + else: + files.extend(path.glob(ext)) + files.sort() documents: list[Document] = [] - for pdf_path in pdf_files: - loader = PyPDFLoader(str(pdf_path)) - pages = loader.load() - - for page in pages: - documents.append( - Document( - text=page.page_content, - metadata={ - "filename": pdf_path.name, - "relative_path": str(pdf_path.relative_to(path)), - "doc_type": "pdf", - "page_number": page.metadata.get("page", 0) + 1, - }, - ) - ) + for file_path in files: + doc_type = EXTENSION_TO_DOC_TYPE[file_path.suffix.lower()] + + if doc_type == "pdf": + documents.extend(_load_pdf(file_path, path)) + else: + documents.extend(_load_single_file(file_path, path, doc_type)) return documents + + +def _load_pdf(file_path: Path, root: Path) -> list[Document]: + """Load a PDF file, returning one Document per page.""" + loader = PyPDFLoader(str(file_path)) + pages = loader.load() + + return [ + Document( + text=page.page_content, + metadata={ + "filename": file_path.name, + "relative_path": str(file_path.relative_to(root)), + "doc_type": "pdf", + "page_number": page.metadata.get("page", 0) + 1, + }, + ) + for page in pages + ] + + +def _load_single_file( + file_path: Path, root: Path, doc_type: str +) -> list[Document]: + """Load a DOCX, TXT, or MD file, returning one Document.""" + if doc_type == "docx": + loader = Docx2txtLoader(str(file_path)) + else: + loader = TextLoader(str(file_path), autodetect_encoding=True) + + pages = loader.load() + text = "\n".join(page.page_content for page in pages) + + return [ + Document( + text=text, + metadata={ + "filename": file_path.name, + "relative_path": str(file_path.relative_to(root)), + "doc_type": doc_type, + "page_number": 1, + }, + ) + ] diff --git a/tests/fixtures/sample.docx b/tests/fixtures/sample.docx new file mode 100644 index 0000000..e657869 Binary files /dev/null and b/tests/fixtures/sample.docx differ diff --git a/tests/fixtures/sample.md b/tests/fixtures/sample.md new file mode 100644 index 0000000..881fe37 --- /dev/null +++ b/tests/fixtures/sample.md @@ -0,0 +1,15 @@ +# Getting Started + +This is the introduction section. + +## Installation + +Run the following command to install. + +### Prerequisites + +You need Python 3.11 or later. + +## Usage + +Import the module and call the main function. diff --git a/tests/fixtures/sample.txt b/tests/fixtures/sample.txt new file mode 100644 index 0000000..22bad60 --- /dev/null +++ b/tests/fixtures/sample.txt @@ -0,0 +1,2 @@ +This is a sample text file for testing. +It contains plain text content across multiple lines. diff --git a/tests/test_answerer.py b/tests/test_answerer.py index 79f0017..5a2fa91 100644 --- a/tests/test_answerer.py +++ b/tests/test_answerer.py @@ -110,3 +110,40 @@ def test_build_source_elements_falls_back_to_filename(): ] elements = build_source_elements(results) assert elements[0]["name"] == "Source: old.pdf | Page 1" + + +def test_source_name_with_section_header(): + """Markdown sources show section header instead of page number.""" + results = [ + SearchResult( + text="Install instructions.", + metadata={ + "filename": "README.md", + "relative_path": "README.md", + "doc_type": "md", + "page_number": 1, + "section_header": "Getting Started > Installation", + }, + distance=0.1, + ), + ] + messages = build_prompt("How to install?", results) + content = " ".join(m["content"] for m in messages) + assert "--- Source: README.md | Section: Getting Started > Installation ---" in content + + elements = build_source_elements(results) + assert elements[0]["name"] == "Source: README.md | Section: Getting Started > Installation" + + +def test_source_name_without_section_header(): + """Non-Markdown sources still show page number (unchanged behaviour).""" + results = [ + SearchResult( + text="Some text.", + metadata={"filename": "report.pdf", "doc_type": "pdf", "page_number": 7}, + distance=0.1, + ), + ] + messages = build_prompt("question", results) + content = " ".join(m["content"] for m in messages) + assert "--- Source: report.pdf | Page 7 ---" in content diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 18c84b8..74a1654 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -67,3 +67,38 @@ def test_chunk_has_doc_hash(): assert "doc_hash" in chunks[0].metadata assert isinstance(chunks[0].metadata["doc_hash"], str) assert len(chunks[0].metadata["doc_hash"]) == 32 # MD5 hex length + + +def test_chunk_markdown_has_section_header(): + """Markdown documents produce chunks with section_header metadata.""" + md_text = "# Introduction\n\nThis is the intro.\n\n## Methods\n\nThese are the methods." + doc = Document( + text=md_text, + metadata={"filename": "README.md", "doc_type": "md", "page_number": 1}, + ) + chunks = chunk_documents([doc], chunk_size=512, chunk_overlap=0) + headers = [c.metadata.get("section_header") for c in chunks] + assert any(h and "Introduction" in h for h in headers) + assert any(h and "Methods" in h for h in headers) + + +def test_chunk_markdown_nested_headers(): + """Nested Markdown headers produce 'H1 > H2' format in section_header.""" + md_text = "# Guide\n\n## Setup\n\n### Prerequisites\n\nYou need Python." + doc = Document( + text=md_text, + metadata={"filename": "guide.md", "doc_type": "md", "page_number": 1}, + ) + chunks = chunk_documents([doc], chunk_size=512, chunk_overlap=0) + headers = [c.metadata.get("section_header", "") for c in chunks] + assert any("Guide" in h and "Setup" in h and "Prerequisites" in h for h in headers) + + +def test_chunk_non_markdown_no_section_header(): + """Non-Markdown documents don't get section_header metadata.""" + doc = Document( + text="Some plain text content.", + metadata={"filename": "notes.txt", "doc_type": "txt", "page_number": 1}, + ) + chunks = chunk_documents([doc], chunk_size=512, chunk_overlap=0) + assert "section_header" not in chunks[0].metadata diff --git a/tests/test_loader.py b/tests/test_loader.py index c8e9978..ce8d284 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -26,10 +26,10 @@ def test_document_metadata_filename(): def test_document_metadata_doc_type(): - """Each Document has doc_type='pdf' in metadata.""" + """Each Document has a valid doc_type in metadata.""" docs = load_folder(FIXTURES) for doc in docs: - assert doc.metadata["doc_type"] == "pdf" + assert doc.metadata["doc_type"] in ("pdf", "docx", "txt", "md") def test_document_metadata_page_number(): @@ -49,7 +49,8 @@ def test_load_folder_empty_dir(tmp_path): def test_document_metadata_relative_path(): """Documents at the root have relative_path equal to filename.""" docs = load_folder(FIXTURES) - for doc in docs: + pdf_docs = [d for d in docs if d.metadata["doc_type"] == "pdf"] + for doc in pdf_docs: assert doc.metadata["relative_path"] == "sample.pdf" @@ -89,3 +90,56 @@ def test_load_folder_non_recursive(tmp_path): docs = load_folder(tmp_path, recursive=False) assert docs == [] + + +def test_load_docx_returns_documents(): + """Loading a folder with a DOCX returns Documents with text.""" + docs = load_folder(FIXTURES) + docx_docs = [d for d in docs if d.metadata["doc_type"] == "docx"] + assert len(docx_docs) > 0 + assert docx_docs[0].text.strip() + + +def test_docx_metadata_doc_type(): + """DOCX documents have doc_type='docx'.""" + docs = load_folder(FIXTURES) + docx_docs = [d for d in docs if d.metadata["filename"] == "sample.docx"] + assert len(docx_docs) > 0 + assert docx_docs[0].metadata["doc_type"] == "docx" + + +def test_docx_metadata_relative_path(): + """DOCX documents have correct relative_path.""" + docs = load_folder(FIXTURES) + docx_docs = [d for d in docs if d.metadata["filename"] == "sample.docx"] + assert docx_docs[0].metadata["relative_path"] == "sample.docx" + + +def test_load_txt_returns_documents(): + """Loading a folder with a TXT file returns Documents with doc_type='txt'.""" + docs = load_folder(FIXTURES) + txt_docs = [d for d in docs if d.metadata["filename"] == "sample.txt"] + assert len(txt_docs) > 0 + assert txt_docs[0].metadata["doc_type"] == "txt" + assert txt_docs[0].text.strip() + + +def test_load_md_returns_documents(): + """Loading a folder with an MD file returns Documents with doc_type='md'.""" + docs = load_folder(FIXTURES) + md_docs = [d for d in docs if d.metadata["filename"] == "sample.md"] + assert len(md_docs) > 0 + assert md_docs[0].metadata["doc_type"] == "md" + assert md_docs[0].text.strip() + + +def test_load_mixed_formats(tmp_path): + """A folder with PDF, DOCX, TXT, and MD loads all formats.""" + import shutil + + for name in ("sample.pdf", "sample.docx", "sample.txt", "sample.md"): + shutil.copy(FIXTURES / name, tmp_path / name) + + docs = load_folder(tmp_path) + doc_types = {d.metadata["doc_type"] for d in docs} + assert doc_types == {"pdf", "docx", "txt", "md"}