diff --git a/src/rag_python/cli.py b/src/rag_python/cli.py index 1c7543b..503efd6 100644 --- a/src/rag_python/cli.py +++ b/src/rag_python/cli.py @@ -154,12 +154,7 @@ def _make_parser() -> argparse.ArgumentParser: "optionally stream tokens and show sources." ), formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=( - "examples:\n" - ' rag-python query "How many days of annual leave?"\n' - " rag-python query \"PTO policy\" --stream -v\n" - ' rag-python query "benefits" --retriever hybrid --metadata-filter \'{"filename": "hr.pdf"}\'' - ), + # Epilog is populated below (after new flags) with expanded Agent-friendly examples. ) q.add_argument( "question", @@ -183,13 +178,58 @@ def _make_parser() -> argparse.ArgumentParser: action="store_true", help="After the answer, print evaluation scores and top source paths", ) + q.add_argument( + "-q", + "--quiet", + action="store_true", + help=( + "Only print the answer text; suppress evaluation and sources trailers. " + "Useful for scripts and CI pipelines." + ), + ) + q.add_argument( + "-f", + "--output-format", + default="text", + choices=["text", "json", "json-pretty"], + metavar="FMT", + help=( + "Output format (default: text). Use 'json' for structured output — " + "designed for calling the CLI as an Agent tool via subprocess with " + "json.loads(). 'json-pretty' adds indentation for human inspection." + ), + ) + q.add_argument( + "--with-sources", + action="store_true", + help=( + "Print the top-5 sources trailer without requiring the full " + "-v/--verbose evaluation section." + ), + ) _add_provider_args(q) _add_search_args(q) - + # Epilog: Agent-friendly flags. Long strings split to stay within line-length=100. + q.epilog = ( + "examples:\n" + ' rag-python query "How many days of annual leave?"\n' + " rag-python query \"PTO policy\" --stream -v\n" + ' rag-python query "benefits" --retriever hybrid' + ' --metadata-filter \'{"filename": "hr.pdf"}\'\n' + " rag-python query \"leave policy\" -q" + " # answer only, script-friendly\n" + ' rag-python query "payroll url" -f json' + ' --quiet # Agent tool: structured JSON\n' + " rag-python query \"refund rule\" --with-sources" + " # sources section, no eval noise\n" + ) docs = sub.add_parser( "docs", help="Show user documentation in the terminal", - description="Print built-in help topics. Full docs: https://github.com/RaghavOG/rag-python/tree/main/docs", + description=( + "Print built-in help topics. " + "Full docs: https://github.com/RaghavOG/rag-python/tree/main/docs" + ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog="topics: " + ", ".join(list_topics()), ) @@ -238,28 +278,61 @@ def main(argv: list[str] | None = None) -> None: retriever=retriever or rag.config.search.retriever, metadata_filter=args.metadata_filter or rag.config.search.metadata_filter, ) + fmt = args.output_format # "text" (default) | "json" | "json-pretty" + + # -------- Highest priority: structured output (Agent / tool use) -------- + # When the caller explicitly requests JSON, we always return the full + # RAGAnswer shape as a single JSON object (stream tokens are discarded + # because JSON output is a final snapshot by definition). + if fmt in ("json", "json-pretty"): + ans = rag.query(question, search=search) + payload = { + "text": ans.text, + "sources": ans.sources, + "evaluation": ans.evaluation, + "retried": ans.retried, + } + if fmt == "json": + print(json.dumps(payload, ensure_ascii=False)) + else: + print(json.dumps(payload, ensure_ascii=False, indent=2)) + return + + # ---------------------- Default: text (human-readable) ---------------------- + quiet = args.quiet + show_sources = args.with_sources or args.verbose + show_evaluation = args.verbose + + def _print_sources_trailer(sources: list[dict]) -> None: + print("\n--- sources ---") + for s in sources[:5]: + meta = s.get("metadata", {}) or {} + print(meta.get("source", ""), "score:", s.get("score")) + if args.stream: stream = rag.query_stream(question, search=search) for token in stream: print(token, end="", flush=True) print() + if quiet: + return result = stream.result - if args.verbose: + if show_evaluation: print("\n--- evaluation ---") print(result.evaluation) - print("\n--- sources ---") - for s in result.sources[:5]: - print(s.get("metadata", {}).get("source", ""), "score:", s.get("score")) + if show_sources: + _print_sources_trailer(result.sources) return ans = rag.query(question, search=search) print(ans.text) - if args.verbose: + if quiet: + return + if show_evaluation: print("\n--- evaluation ---") print(ans.evaluation) - print("\n--- sources ---") - for s in ans.sources[:5]: - print(s.get("metadata", {}).get("source", ""), "score:", s.get("score")) + if show_sources: + _print_sources_trailer(ans.sources) if __name__ == "__main__": diff --git a/tests/test_cli_query_agent_output.py b/tests/test_cli_query_agent_output.py new file mode 100644 index 0000000..f552593 --- /dev/null +++ b/tests/test_cli_query_agent_output.py @@ -0,0 +1,272 @@ +"""Tests for CLI query Agent-friendly output (PR #5 feature). + +Features under test: + - -q/--quiet flag: print ONLY the answer text (no evaluation/sources trailers) + - --output-format / -f: text (default compat) | json (compact) | json-pretty (indented) + → JSON output exposes the full RAGAnswer structure (text, sources, evaluation, retried) + → so Agent frameworks can subprocess.call the CLI and json.loads the result + - --with-sources: print sources section independently, without requiring full --verbose + +RED phase: ALL tests FAIL because cli.py does not implement these flags yet. +""" +from __future__ import annotations + +import io +import json +import subprocess +import sys +from unittest.mock import MagicMock, patch + +# ---- Resolve path to src/ so tests run without pip-install ---- +from pathlib import Path +_SRC = Path(__file__).resolve().parent.parent / "src" +if str(_SRC) not in sys.path: + sys.path.insert(0, str(_SRC)) + +from rag_python.client import RAGAnswer # noqa: E402 (import after sys.path fix) +from rag_python.options import SearchConfig # noqa: E402 (real dataclass for replace()) + + +def _stub_answer() -> RAGAnswer: + """A fixed RAGAnswer object the stub RAG.query() returns.""" + return RAGAnswer( + text="You have 15 days of annual leave.", + sources=[ + {"metadata": {"source": "hr_policy.pdf", "page": 3}, "score": 0.87}, + {"metadata": {"source": "handbook.md"}, "score": 0.61}, + {"metadata": {"source": "slack_thread.txt"}, "score": 0.55}, + {"metadata": {"source": "wiki.docx"}, "score": 0.48}, + {"metadata": {"source": "extra_a.pdf"}, "score": 0.42}, + {"metadata": {"source": "extra_b.pdf"}, "score": 0.40}, # beyond top-5 cutoff + ], + evaluation={ + "faithfulness": 0.92, + "relevancy": 0.88, + "chunks_retrieved": 5, + }, + retried=False, + ) + + +def _run_main_query(argv): + """Run cli.main(argv) with _build_rag mocked to return a stub RAG. + + The stub RAG: + - has a .config.search attribute (dataclass-like, supports replace() fields) + - .query(question, search=...) returns a fixed RAGAnswer + - .query_stream(question, search=...) returns a stub stream object + + Returns (stdout, stderr, stub_rag). + """ + fake_rag = MagicMock() + # Use a REAL SearchConfig dataclass so dataclasses.replace() works inside cli.py + # Fields match cli.py's replace() calls: retriever, metadata_filter + fake_rag.config.search = SearchConfig(retriever="multi_query", metadata_filter=None) + # Default query() return value + fake_rag.query.return_value = _stub_answer() + + # --- stub for query_stream path (used by --stream tests) --- + class _StubStream: + def __init__(self, tokens, answer: RAGAnswer): + self._tokens = tokens + self.result = answer + def __iter__(self): + return iter(self._tokens) + fake_rag.query_stream.return_value = _StubStream( + tokens=["Hello", " ", "world", "."], + answer=_stub_answer(), + ) + + with patch("rag_python.cli._build_rag", return_value=fake_rag): + buf_out = io.StringIO() + buf_err = io.StringIO() + old_out, old_err = sys.stdout, sys.stderr + try: + sys.stdout, sys.stderr = buf_out, buf_err + from rag_python.cli import main + main(argv) + except SystemExit as e: + if e.code not in (0, None): + raise + finally: + sys.stdout, sys.stderr = old_out, old_err + return buf_out.getvalue(), buf_err.getvalue(), fake_rag + + +# ============================================================ +# Test 1 – query --help lists --quiet / --output-format / --with-sources +# ============================================================ +def test_query_help_lists_new_flags(): + result = subprocess.run( + [sys.executable, "-m", "rag_python.cli", "query", "--help"], + capture_output=True, + text=True, + check=False, + env={**dict(__import__("os").environ), "PYTHONPATH": str(_SRC)}, + ) + assert result.returncode == 0, f"query --help failed: {result.stderr}" + assert "--quiet" in result.stdout, ( + f"Missing '--quiet' in query --help. Stdout:\n{result.stdout}" + ) + assert "--output-format" in result.stdout or "-f" in result.stdout, ( + f"Missing '--output-format/-f' in query --help. Stdout:\n{result.stdout}" + ) + assert "--with-sources" in result.stdout, ( + f"Missing '--with-sources' in query --help. Stdout:\n{result.stdout}" + ) + + +# ============================================================ +# Test 2 – --quiet suppresses evaluation/sources trailers (only text answer) +# ============================================================ +def test_query_quiet_only_prints_answer_text(): + stdout, stderr, _ = _run_main_query( + ["query", "--quiet", "How much annual leave do I have?"] + ) + # stderr must be empty (no argparse errors) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + # Answer text must appear exactly as the trimmed stdout + assert stdout.strip() == "You have 15 days of annual leave.", ( + f"With --quiet expected ONLY the answer text on stdout. Got:\n{stdout!r}" + ) + # Quiet mode: absolutely NO evaluation/sources headers + assert "--- evaluation ---" not in stdout, "Evaluation header leaked in --quiet mode" + assert "--- sources ---" not in stdout, "Sources header leaked in --quiet mode" + + +# ============================================================ +# Test 3 – short alias -q behaves identical to --quiet +# ============================================================ +def test_query_short_q_alias_matches_quiet(): + stdout_q, stderr_q, _ = _run_main_query( + ["query", "-q", "leave policy?"] + ) + stdout_long, stderr_long, _ = _run_main_query( + ["query", "--quiet", "leave policy?"] + ) + # No crashes + assert stderr_q.strip() == "", f"-q produced stderr: {stderr_q!r}" + assert stderr_long.strip() == "", f"--quiet produced stderr: {stderr_long!r}" + # Same output + assert stdout_q == stdout_long, ( + f"Short alias -q does not match --quiet.\n-q: {stdout_q!r}\n--quiet: {stdout_long!r}" + ) + + +# ============================================================ +# Test 4 – --output-format json produces VALID, parseable JSON +# with the 4 RAGAnswer keys, correct values, compact (no trailing '\n' count rule) +# ============================================================ +def test_query_output_format_json_is_valid_and_structured(): + stdout, stderr, _ = _run_main_query( + ["query", "--output-format", "json", "my question"] + ) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + # Must be parseable JSON — this is the CRITICAL assertion for Agent Tool use + try: + data = json.loads(stdout) + except json.JSONDecodeError as e: + raise AssertionError( + f"--output-format json produced non-JSON stdout: {e}\nRaw:\n{stdout!r}" + ) from None + # 4 keys from RAGAnswer shape + for required_key in ("text", "sources", "evaluation", "retried"): + assert required_key in data, f"JSON missing key: {required_key!r}. Keys: {list(data)}" + expected = _stub_answer() + assert data["text"] == expected.text, f"text mismatch: {data['text']!r} != {expected.text!r}" + assert isinstance(data["sources"], list), "sources must be a list" + assert len(data["sources"]) == len(expected.sources), ( + f"JSON should expose ALL sources (not truncated to top-5). " + f"Got {len(data['sources'])}, expected {len(expected.sources)}" + ) + assert data["evaluation"] == expected.evaluation, "evaluation dict mismatch" + assert data["retried"] is False, "retried field mismatch" + + +# ============================================================ +# Test 5 – --output-format json-pretty: valid JSON AND pretty-printed +# ============================================================ +def test_query_output_format_json_pretty_has_indentation(): + stdout, stderr, _ = _run_main_query( + ["query", "-f", "json-pretty", "q"] + ) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + # Parseable JSON + try: + data = json.loads(stdout) + except json.JSONDecodeError as e: + raise AssertionError( + f"json-pretty produced non-JSON stdout: {e}\nRaw:\n{stdout!r}" + ) from None + assert set(("text", "sources", "evaluation", "retried")).issubset(data.keys()) + # Pretty = contains newlines inside the JSON object (not just trailing) + non_trailing_newlines = len(stdout.rstrip("\n").split("\n")) + assert non_trailing_newlines >= 4, ( + f"json-pretty should be multi-line indentation. " + f"Got {non_trailing_newlines} lines. Raw:\n{stdout!r}" + ) + + +# ============================================================ +# Test 6 – --with-sources prints sources section WITHOUT requiring --verbose +# ============================================================ +def test_query_with_sources_prints_sources_independently(): + stdout, stderr, _ = _run_main_query( + ["query", "--with-sources", "leave policy"] + ) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + # Answer text still present + assert "15 days of annual leave" in stdout, "Answer missing in stdout" + # Sources header present (without needing -v / --verbose) + assert "--- sources ---" in stdout, ( + f"--with-sources did not trigger sources section. Stdout:\n{stdout!r}" + ) + # WITHOUT --verbose → evaluation section must NOT appear + assert "--- evaluation ---" not in stdout, ( + "--with-sources should NOT imply verbose evaluation output" + ) + + +# ============================================================ +# Test 7 – Default (no new flags) behaviour preserved: backwards-compat regression +# ============================================================ +def test_query_default_no_flags_preserves_original_behavior(): + """Original behaviour: + → print(ans.text) alone (no sources, no evaluation) + → unless -v/--verbose is passed which adds eval + sources trailers + """ + # --- sub-case A: no flags → only answer text on stdout --- + stdout, stderr, _ = _run_main_query( + ["query", "leave?"] + ) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + assert stdout.strip() == "You have 15 days of annual leave.", ( + f"Default no-flags stdout changed. Got:\n{stdout!r}" + ) + assert "--- evaluation ---" not in stdout + assert "--- sources ---" not in stdout + + # --- sub-case B: -v (original verbose) → adds BOTH eval + sources --- + stdout_v, stderr_v, _ = _run_main_query( + ["query", "-v", "leave?"] + ) + assert stderr_v.strip() == "", f"Unexpected stderr on -v: {stderr_v!r}" + assert "--- evaluation ---" in stdout_v, "-v verbose lost evaluation section" + assert "--- sources ---" in stdout_v, "-v verbose lost sources section" + + +# ============================================================ +# Test 8 – --stream mode + --quiet: tokens stream, eval/sources suppressed +# ============================================================ +def test_query_stream_quiet_streams_tokens_no_tail(): + stdout, stderr, _ = _run_main_query( + ["query", "--stream", "--quiet", "q"] + ) + assert stderr.strip() == "", f"Unexpected stderr: {stderr!r}" + # Stream path → tokens are printed (we stubbed "Hello world.") + assert "Hello" in stdout and "world" in stdout, ( + f"Streaming tokens not printed. Stdout:\n{stdout!r}" + ) + # Quiet → NO eval/sources tail headers + assert "--- evaluation ---" not in stdout + assert "--- sources ---" not in stdout