diff --git a/README.md b/README.md index fee4b86..8096e7b 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?" +# add --json for machine-readable answer and 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..6f043a0 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="Emit the answer and sources as JSON", + ) return p @@ -47,8 +53,26 @@ 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: + answer = answerer.answer(args.question, hits) + if args.json: + print( + json.dumps( + { + "answer": answer, + "sources": [ + { + "title": note.title, + "file": note.path.name, + "score": score, + } + for note, score in hits + ], + } + ) + ) + else: + print(answer) + if hits and not args.json: print("\nSources:") for note, score in hits: print(f" - {note.title} ({note.path.name}, score={score:.3f})") diff --git a/tests/test_notes_rag.py b/tests/test_notes_rag.py index 925a97d..7a210de 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,15 @@ def test_mock_answer_cites(capsys): assert "Sources:" in out +def test_mock_json_output(capsys): + code = main(["--mock", "--json", "--notes", str(NOTES), "What is a git remote?"]) + out = capsys.readouterr().out + assert code == 0 + result = json.loads(out) + assert result["answer"].startswith("[mock]") + assert result["sources"] + + def test_empty_notes_dir(tmp_path, capsys): empty = tmp_path / "empty" empty.mkdir()