diff --git a/README.md b/README.md index 4b09611..7f01015 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ cd notes-rag pip install -e ".[dev]" pytest python -m notes_rag --mock "What is a git remote?" +python -m notes_rag --mock --json "What is a git remote?" # machine-readable answer + sources ``` You should see a `[mock]` answer plus a **Sources** list. That means retrieval works before you spend a token. diff --git a/src/notes_rag/cli.py b/src/notes_rag/cli.py index c0d1c34..99610fb 100644 --- a/src/notes_rag/cli.py +++ b/src/notes_rag/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse +import json import sys from pathlib import Path @@ -34,6 +35,11 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Answer from retrieved notes without calling an API", ) + p.add_argument( + "--json", + action="store_true", + help="Print a single JSON object with answer and sources", + ) return p @@ -47,11 +53,19 @@ def main(argv: list[str] | None = None) -> int: hits = retrieve(notes, args.question, k=args.k) answerer = Answerer.from_env(mock=args.mock) try: - print(answerer.answer(args.question, hits)) - if hits: - print("\nSources:") - for note, score in hits: - print(f" - {note.title} ({note.path.name}, score={score:.3f})") + answer = answerer.answer(args.question, hits) + if args.json: + sources = [ + {"title": note.title, "file": note.path.name, "score": score} + for note, score in hits + ] + print(json.dumps({"answer": answer, "sources": sources})) + else: + print(answer) + if hits: + print("\nSources:") + for note, score in hits: + print(f" - {note.title} ({note.path.name}, score={score:.3f})") except Exception as exc: # noqa: BLE001 print(f"error: {exc}", file=sys.stderr) return 1 diff --git a/tests/test_notes_rag.py b/tests/test_notes_rag.py index 925a97d..745784e 100644 --- a/tests/test_notes_rag.py +++ b/tests/test_notes_rag.py @@ -1,3 +1,4 @@ +import json from pathlib import Path from notes_rag.answer import Answerer @@ -37,6 +38,16 @@ def test_mock_answer_cites(capsys): assert "Sources:" in out +def test_mock_answer_json_output(capsys): + code = main(["--mock", "--json", "--notes", str(NOTES), "What is a git remote?"]) + out = capsys.readouterr().out + data = json.loads(out) + assert code == 0 + assert "[mock]" in data["answer"] + assert data["sources"] + assert {"title", "file", "score"}.issubset(data["sources"][0]) + + def test_empty_notes_dir(tmp_path, capsys): empty = tmp_path / "empty" empty.mkdir()