diff --git a/nemo_retriever/src/nemo_retriever/common/input_files.py b/nemo_retriever/src/nemo_retriever/common/input_files.py index 5f015dab41..99f9987bef 100644 --- a/nemo_retriever/src/nemo_retriever/common/input_files.py +++ b/nemo_retriever/src/nemo_retriever/common/input_files.py @@ -12,6 +12,9 @@ "*.docx", "*.pptx", "*.txt", + "*.md", + "*.json", + "*.sh", "*.html", "*.jpg", "*.jpeg", @@ -28,7 +31,7 @@ "*.mkv", ), "pdf": ("*.pdf",), - "txt": ("*.txt",), + "txt": ("*.txt", "*.md", "*.json", "*.sh"), "html": ("*.html",), "doc": ("*.docx", "*.pptx"), "image": ("*.jpg", "*.jpeg", "*.png", "*.tiff", "*.tif", "*.bmp", "*.svg"), @@ -130,7 +133,9 @@ def resolve_input_files(input_path: Path, input_type: str) -> list[Path]: if not path.exists(): return [] - files: list[Path] = [] - for pattern in INPUT_TYPE_PATTERNS.get(input_type, INPUT_TYPE_PATTERNS["pdf"]): - files.extend(match for match in path.rglob(pattern) if match.is_file()) - return sorted(set(files)) + allowed_extensions = ( + AUTO_INPUT_EXTENSIONS + if input_type == "auto" + else INPUT_TYPE_EXTENSIONS.get(input_type, INPUT_TYPE_EXTENSIONS["pdf"]) + ) + return sorted(match for match in path.rglob("*") if match.is_file() and match.suffix.lower() in allowed_extensions) diff --git a/nemo_retriever/src/nemo_retriever/service/utils/file_type.py b/nemo_retriever/src/nemo_retriever/service/utils/file_type.py index 07d20e96a7..367ce57b51 100644 --- a/nemo_retriever/src/nemo_retriever/service/utils/file_type.py +++ b/nemo_retriever/src/nemo_retriever/service/utils/file_type.py @@ -51,6 +51,9 @@ class FileClassifier: ".pptx": (FileCategory.DOCUMENT, "application/vnd.openxmlformats-officedocument.presentationml.presentation"), # Plain text ".txt": (FileCategory.TEXT, "text/plain"), + ".md": (FileCategory.TEXT, "text/plain"), + ".json": (FileCategory.TEXT, "text/plain"), + ".sh": (FileCategory.TEXT, "text/plain"), # Web / markup ".html": (FileCategory.HTML, "text/html"), # Image diff --git a/nemo_retriever/tests/service/utils/test_file_type.py b/nemo_retriever/tests/service/utils/test_file_type.py new file mode 100644 index 0000000000..eb8fae6741 --- /dev/null +++ b/nemo_retriever/tests/service/utils/test_file_type.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from io import BytesIO + +import pytest +from fastapi import UploadFile + +from nemo_retriever.service.utils.file_type import FileCategory, FileClassifier + + +@pytest.mark.parametrize("filename", ["README.MD", "payload.json", "setup.sh"]) +def test_documented_plain_text_extensions_are_classified_like_txt(filename: str) -> None: + upload = UploadFile(filename=filename, file=BytesIO(b"plain text content")) + + classification = FileClassifier.classify(upload) + + assert classification.category == FileCategory.TEXT + assert classification.content_type == "text/plain" diff --git a/nemo_retriever/tests/test_graph_pipeline_cli.py b/nemo_retriever/tests/test_graph_pipeline_cli.py index 47dba005f5..4151e5647b 100644 --- a/nemo_retriever/tests/test_graph_pipeline_cli.py +++ b/nemo_retriever/tests/test_graph_pipeline_cli.py @@ -129,7 +129,12 @@ def test_resolve_input_file_patterns_recurses_for_directory_inputs(tmp_path) -> doc_patterns = resolve_input_patterns(dataset_dir, "doc") assert pdf_patterns == [str(dataset_dir / "**" / "*.pdf")] - assert txt_patterns == [str(dataset_dir / "**" / "*.txt")] + assert txt_patterns == [ + str(dataset_dir / "**" / "*.txt"), + str(dataset_dir / "**" / "*.md"), + str(dataset_dir / "**" / "*.json"), + str(dataset_dir / "**" / "*.sh"), + ] assert doc_patterns == [str(dataset_dir / "**" / "*.docx"), str(dataset_dir / "**" / "*.pptx")] diff --git a/nemo_retriever/tests/test_ingest_interface.py b/nemo_retriever/tests/test_ingest_interface.py index d74f2705d9..00e5b4cc47 100644 --- a/nemo_retriever/tests/test_ingest_interface.py +++ b/nemo_retriever/tests/test_ingest_interface.py @@ -1,3 +1,5 @@ +from io import BytesIO +from pathlib import Path from types import SimpleNamespace import pandas as pd @@ -269,6 +271,40 @@ def test_extract_default_rejects_unknown_input_type(tmp_path) -> None: ingestor.ingest() +def test_extract_default_treats_markdown_as_plain_text(tmp_path) -> None: + document = tmp_path / "README.md" + document.write_text("# Heading\n\nBody text\n", encoding="utf-8") + + result = GraphIngestor(run_mode="inprocess", show_progress=False).files([str(document)]).extract().ingest() + + assert result["text"].tolist() == ["# Heading\n\nBody text\n"] + assert result["path"].tolist() == [str(document.resolve())] + + +def test_extract_txt_accepts_json_as_plain_text(tmp_path) -> None: + document = tmp_path / "payload.json" + document.write_text('{"message": "hello"}\n', encoding="utf-8") + + result = GraphIngestor(run_mode="inprocess", show_progress=False).files([str(document)]).extract_txt().ingest() + + assert result["text"].tolist() == ['{"message": "hello"}\n'] + assert result["path"].tolist() == [str(document.resolve())] + + +def test_extract_default_accepts_shell_script_buffer_as_plain_text() -> None: + content = b"#!/bin/sh\necho hello\n" + + result = ( + GraphIngestor(run_mode="inprocess", show_progress=False) + .buffers(("setup.sh", BytesIO(content))) + .extract() + .ingest() + ) + + assert result["text"].tolist() == [content.decode()] + assert result["path"].tolist() == [str(Path("setup.sh").resolve())] + + def test_typed_shortcuts_preserve_legacy_no_default_chunking() -> None: """Typed shortcuts (extract_audio, extract_txt, ...) must NOT enable default split_config chunking. Default-ON is reserved for the unified .extract() diff --git a/nemo_retriever/tests/test_ingest_manifest.py b/nemo_retriever/tests/test_ingest_manifest.py index 2e74bbbde0..f754792bdc 100644 --- a/nemo_retriever/tests/test_ingest_manifest.py +++ b/nemo_retriever/tests/test_ingest_manifest.py @@ -110,6 +110,18 @@ def test_manifest_planner_mixed_inputs_use_stable_family_order(tmp_path) -> None assert [branch.family for branch in branches] == ["pdf", "image", "txt"] +@pytest.mark.parametrize("suffix", [".md", ".json", ".sh"]) +def test_manifest_planner_routes_documented_plain_text_extensions_to_text(tmp_path, suffix) -> None: + document = tmp_path / f"document{suffix}" + document.write_text("plain text content", encoding="utf-8") + + branches = plan_extraction_branches(build_input_manifest([str(document)])) + + assert [(branch.family, branch.extraction_mode, branch.input_paths) for branch in branches] == [ + ("txt", "text", (str(document),)), + ] + + def test_manifest_branch_specs_resolve_default_params(monkeypatch, tmp_path) -> None: audio = tmp_path / "clip.wav" video = tmp_path / "scene.mp4" diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index 44ba7dcf03..23d49b4929 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -1022,6 +1022,22 @@ def test_root_ingest_routes_text_inputs_by_default_to_auto_planner(monkeypatch, assert isinstance(fake_ingestor.extract.call_args.kwargs["text_params"], TextChunkParams) +@pytest.mark.parametrize("suffix", [".md", ".json", ".sh"]) +def test_root_ingest_treats_documented_plain_text_extensions_as_text(monkeypatch, tmp_path, suffix) -> None: + fake_ingestor = _make_fake_ingestor() + document = tmp_path / f"notes{suffix}" + document.write_text("plain text content", encoding="utf-8") + + monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) + + result = RUNNER.invoke(cli_main.app, ["ingest", str(document)]) + + assert result.exit_code == 0 + assert fake_ingestor.files.call_args.args == ([str(document)],) + assert isinstance(fake_ingestor.extract.call_args.args[0], ExtractParams) + assert isinstance(fake_ingestor.extract.call_args.kwargs["text_params"], TextChunkParams) + + def test_root_ingest_help_defaults_to_local_workflow(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(typer_rich_utils, "MAX_WIDTH", 200) monkeypatch.setattr(typer_rich_utils, "FORCE_TERMINAL", False) @@ -1474,19 +1490,23 @@ def test_root_ingest_auto_mixed_directory_uses_auto_extraction(monkeypatch, tmp_ assert isinstance(fake_ingestor.extract.call_args.kwargs["text_params"], TextChunkParams) -def test_root_ingest_text_formats_directory_skips_markdown(monkeypatch, tmp_path) -> None: +def test_root_ingest_text_formats_directory_includes_documented_plain_text_extensions(monkeypatch, tmp_path) -> None: fake_ingestor = _make_fake_ingestor() dataset = tmp_path / "data" dataset.mkdir() html = dataset / "architecture.html" text = dataset / "api_changelog.txt" markdown = dataset / "aurora_README.md" + json_document = dataset / "metadata.json" + shell_script = dataset / "setup.sh" html.write_text("

Architecture

", encoding="utf-8") text.write_text("API changelog", encoding="utf-8") - markdown.write_text("# Unsupported markdown", encoding="utf-8") + markdown.write_text("# Aurora", encoding="utf-8") + json_document.write_text('{"project": "aurora"}', encoding="utf-8") + shell_script.write_text("#!/bin/sh\necho aurora\n", encoding="utf-8") monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setattr(ingest_execution, "_count_lancedb_rows", lambda *_, **__: 2) + monkeypatch.setattr(ingest_execution, "_count_lancedb_rows", lambda *_, **__: 5) result = RUNNER.invoke(cli_main.app, ["ingest", str(dataset)]) @@ -1494,12 +1514,44 @@ def test_root_ingest_text_formats_directory_skips_markdown(monkeypatch, tmp_path assert set(fake_ingestor.files.call_args.args[0]) == { str(html.resolve()), str(text.resolve()), + str(markdown.resolve()), + str(json_document.resolve()), + str(shell_script.resolve()), } - assert str(markdown.resolve()) not in fake_ingestor.files.call_args.args[0] extract_kwargs = fake_ingestor.extract.call_args.kwargs assert isinstance(extract_kwargs["text_params"], TextChunkParams) assert isinstance(extract_kwargs["html_params"], HtmlChunkParams) - assert "Ingested 2 file(s) → 2 row(s)" in result.output + assert "Ingested 5 file(s) → 5 row(s)" in result.output + + +def test_root_ingest_directory_discovers_text_extensions_case_insensitively(monkeypatch, tmp_path) -> None: + fake_ingestor = _make_fake_ingestor() + document = tmp_path / "README.MD" + document.write_text("# Heading\n", encoding="utf-8") + + monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) + + result = RUNNER.invoke(cli_main.app, ["ingest", str(tmp_path)]) + + assert result.exit_code == 0 + assert fake_ingestor.files.call_args.args == ([str(document.resolve())],) + assert isinstance(fake_ingestor.extract.call_args.kwargs["text_params"], TextChunkParams) + + +def test_root_ingest_expands_documented_plain_text_glob(monkeypatch, tmp_path) -> None: + fake_ingestor = _make_fake_ingestor() + nested = tmp_path / "nested" + nested.mkdir() + script = nested / "setup.sh" + script.write_text("#!/bin/sh\necho hello\n", encoding="utf-8") + + monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) + + result = RUNNER.invoke(cli_main.app, ["ingest", str(tmp_path / "**" / "*.sh")]) + + assert result.exit_code == 0 + assert fake_ingestor.files.call_args.args == ([str(script)],) + assert isinstance(fake_ingestor.extract.call_args.kwargs["text_params"], TextChunkParams) def test_root_ingest_reports_os_errors(monkeypatch) -> None: diff --git a/nemo_retriever/tests/test_service_pipeline_spec.py b/nemo_retriever/tests/test_service_pipeline_spec.py index f6bb46eb3a..eca2567f2b 100644 --- a/nemo_retriever/tests/test_service_pipeline_spec.py +++ b/nemo_retriever/tests/test_service_pipeline_spec.py @@ -483,6 +483,9 @@ def test_build_graph_ingestor_attaches_asr_params_for_explicit_audio_mode() -> N ("filename", "expected"), [ ("notes.txt", "text"), + ("README.md", "text"), + ("payload.json", "text"), + ("setup.sh", "text"), ("page.html", "html"), ("report.pdf", "pdf"), ("diagram.png", "image"), @@ -513,8 +516,8 @@ def test_build_graph_ingestor_uses_typed_txt_html_shortcuts() -> None: spec = {"extraction_mode": "auto", "stage_order": ["extract"]} txt_ingestor, txt_mode, _ = _build_graph_ingestor_from_spec( - "notes.txt", - b"The quick brown fox", + "README.md", + b"# The quick brown fox", base_extract, None, spec, diff --git a/skills/nemo-retriever/SKILL.md b/skills/nemo-retriever/SKILL.md index b8ed063cff..690810667c 100644 --- a/skills/nemo-retriever/SKILL.md +++ b/skills/nemo-retriever/SKILL.md @@ -1,6 +1,6 @@ --- name: nemo-retriever -description: "Use when the user wants to search, query, extract, transcribe, describe, quote, filter, or aggregate across documents — PDFs, scanned forms / images (`.jpg` `.png` `.tiff`), Office (`.docx` `.pptx`), text (`.html` `.txt`), audio (`.mp3` `.wav` `.m4a`), or video (`.mp4` `.mov`). Prefer this over native Read / Grep for multi-file or non-PDF corpora. Not for: editing files, web browsing, single-file plain-text lookups, fine-tuning." +description: "Use when the user wants to search, query, extract, transcribe, describe, quote, filter, or aggregate across documents — PDFs, scanned forms / images (`.jpg` `.png` `.tiff`), Office (`.docx` `.pptx`), text (`.html` `.txt` `.md` `.json` `.sh`), audio (`.mp3` `.wav` `.m4a`), or video (`.mp4` `.mov`). Prefer this over native Read / Grep for multi-file or non-PDF corpora. Not for: editing files, web browsing, single-file plain-text lookups, fine-tuning." license: Apache-2.0 allowed-tools: Bash Write Read --- @@ -9,7 +9,7 @@ allowed-tools: Bash Write Read The `retriever` CLI indexes a folder of PDFs into LanceDB (`retriever ingest`) and serves vector search over it (`retriever query`). For any task about searching/answering questions across a folder of PDFs, use this CLI — do not write a custom RAG. -**Beyond PDFs and beyond semantic search.** `retriever ingest` also handles images, Office, HTML, TXT, audio, and video — see `references/setup.md` for the per-format recipe and `references/install.md` for the install extras (`[multimedia]`, libreoffice, ffmpeg). The query turn is two retrieval passes — see **§Query turn** below (inline, no reference read needed); `references/cli/query.md` holds only the fallback detail (exact-term, chart text-extract, compose-reply). Don't fall back to native Read/Grep/Python on non-PDF inputs. +**Beyond PDFs and beyond semantic search.** `retriever ingest` also handles images, Office, HTML, plain-text formats, audio, and video — see `references/setup.md` for the per-format recipe and `references/install.md` for the install extras (`[multimedia]`, libreoffice, ffmpeg). The query turn is two retrieval passes — see **§Query turn** below (inline, no reference read needed); `references/cli/query.md` holds only the fallback detail (exact-term, chart text-extract, compose-reply). Don't fall back to native Read/Grep/Python on non-PDF inputs. ## Install (if `retriever` is missing) diff --git a/skills/nemo-retriever/references/cli/ingest.md b/skills/nemo-retriever/references/cli/ingest.md index 6049a3c82e..b663151b17 100644 --- a/skills/nemo-retriever/references/cli/ingest.md +++ b/skills/nemo-retriever/references/cli/ingest.md @@ -75,8 +75,9 @@ Write to a custom DB / table: - Positional `DOCUMENTS...` is required and repeatable. - Values may be file paths, directories, or shell globs. - Supported input families are detected automatically from extensions: - `pdf`, `docx`, `pptx`, `txt`, `html`, `jpg`, `jpeg`, `png`, `tiff`, `tif`, - `bmp`, `svg`, `mp3`, `wav`, `m4a`, `mp4`, `mov`, and `mkv`. + `pdf`, `docx`, `pptx`, `txt`, `md`, `json`, `sh`, `html`, `jpg`, `jpeg`, + `png`, `tiff`, `tif`, `bmp`, `svg`, `mp3`, `wav`, `m4a`, `mp4`, `mov`, and + `mkv`. Markdown, JSON, and shell scripts are treated as plain text. ## Outputs diff --git a/skills/nemo-retriever/references/setup.md b/skills/nemo-retriever/references/setup.md index aaa9fc6490..3e5f454b37 100644 --- a/skills/nemo-retriever/references/setup.md +++ b/skills/nemo-retriever/references/setup.md @@ -62,8 +62,9 @@ Install extras for non-PDF media live in `references/install.md` under For mixed-script docs such as bilingual contracts or multilingual forms, use `--ocr-lang multi`. Chart understanding runs inline; no separate call is needed. -**HTML / TXT** - ingest even though `Read` could work; chunking and citation -metadata matter: +**HTML / text** (`.html`, `.txt`, `.md`, `.json`, `.sh`) - ingest even though +`Read` could work; chunking and citation metadata matter. Markdown, JSON, and +shell scripts are treated as plain text: ```bash /bin/retriever ingest ./docs/ diff --git a/skills/nemo-retriever/references/troubleshooting.md b/skills/nemo-retriever/references/troubleshooting.md index 00095e612a..9bfbc8cfb0 100644 --- a/skills/nemo-retriever/references/troubleshooting.md +++ b/skills/nemo-retriever/references/troubleshooting.md @@ -27,10 +27,11 @@ For an unlisted subcommand: `/bin/retriever --help` ## Unsupported file types `retriever ingest` auto-detects supported input types from file extensions. It -supports `.pdf`, `.docx`, `.pptx`, `.txt`, `.html`, `.jpg`, `.jpeg`, `.png`, -`.tiff`, `.tif`, `.bmp`, `.svg`, `.mp3`, `.wav`, `.m4a`, `.mp4`, `.mov`, and -`.mkv`. Treat other extensions such as `.flac`, `.rtf`, `.eml`, `.py`, `.jsonl`, -and `.zip` as setup issues. Before ingest, inventory: +supports `.pdf`, `.docx`, `.pptx`, `.txt`, `.md`, `.json`, `.sh`, `.html`, +`.jpg`, `.jpeg`, `.png`, `.tiff`, `.tif`, `.bmp`, `.svg`, `.mp3`, `.wav`, +`.m4a`, `.mp4`, `.mov`, and `.mkv`. Markdown, JSON, and shell scripts use the +plain-text extraction path. Treat other extensions such as `.flac`, `.rtf`, +`.eml`, `.py`, `.jsonl`, and `.zip` as setup issues. Before ingest, inventory: ```bash find -type f -name '*.*' | sed 's/.*\.//' | sort -u diff --git a/skills/nemo-retriever/skill-card.md b/skills/nemo-retriever/skill-card.md index 7216002994..e666e53bb3 100644 --- a/skills/nemo-retriever/skill-card.md +++ b/skills/nemo-retriever/skill-card.md @@ -1,5 +1,5 @@ ## Description:
-Use when the user wants to search, query, extract, transcribe, describe, quote, filter, or aggregate across documents — PDFs, scanned forms / images (.jpg .png .tiff), Office (.docx .pptx), text (.html .txt), audio (.mp3 .wav .m4a), or video (.mp4 .mov).
+Use when the user wants to search, query, extract, transcribe, describe, quote, filter, or aggregate across documents — PDFs, scanned forms / images (.jpg .png .tiff), Office (.docx .pptx), text (.html .txt .md .json .sh), audio (.mp3 .wav .m4a), or video (.mp4 .mov).
This skill is ready for commercial/non-commercial use.