diff --git a/docs/docs/getting_started/starter_tools/rag_cli.md b/docs/docs/getting_started/starter_tools/rag_cli.md index 5b378eea1d6be25b013b83bc034e2453f342fb64..d86d44d1522653bd501a9b196f9caf2b1d5be09e 100644 --- a/docs/docs/getting_started/starter_tools/rag_cli.md +++ b/docs/docs/getting_started/starter_tools/rag_cli.md @@ -22,14 +22,15 @@ After that, you can start using the tool: ```shell $ llamaindex-cli rag -h -usage: llamaindex-cli rag [-h] [-q QUESTION] [-f FILES] [-c] [-v] [--clear] [--create-llama] +usage: llamaindex-cli rag [-h] [-q QUESTION] [-f FILES [FILES ...]] [-c] [-v] [--clear] [--create-llama] options: -h, --help show this help message and exit -q QUESTION, --question QUESTION The question you want to ask. - -f FILES, --files FILES - The name of the file or directory you want to ask a question about,such as "file.pdf". + -f, --files FILES [FILES ...] + The name of the file(s) or directory you want to ask a question about,such + as "file.pdf". Supports globs like "*.py". -c, --chat If flag is present, opens a chat REPL. -v, --verbose Whether to print out verbose information during execution. --clear Clears out all currently embedded data. diff --git a/llama-index-cli/llama_index/cli/rag/base.py b/llama-index-cli/llama_index/cli/rag/base.py index 61259c10681d7dada0e7891d1e2b620ea2bbadec..105362a9e2c8709c29b51c1f1c3e567ef6a6ddc8 100644 --- a/llama-index-cli/llama_index/cli/rag/base.py +++ b/llama-index-cli/llama_index/cli/rag/base.py @@ -4,7 +4,7 @@ import shutil from argparse import ArgumentParser from glob import iglob from pathlib import Path -from typing import Any, Callable, Dict, Optional, Union, cast +from typing import Any, Callable, Dict, List, Optional, Union, cast from llama_index.core import ( Settings, @@ -176,7 +176,7 @@ class RagCLI(BaseModel): async def handle_cli( self, - files: Optional[str] = None, + files: Optional[List[str]] = None, question: Optional[str] = None, chat: bool = False, verbose: bool = False, @@ -205,8 +205,11 @@ class RagCLI(BaseModel): if self.verbose: print("Saving/Loading from persist_dir: ", self.persist_dir) if files is not None: + expanded_files = [] + for pattern in files: + expanded_files.extend(iglob(pattern, recursive=True)) documents = [] - for _file in iglob(files, recursive=True): + for _file in expanded_files: _file = os.path.abspath(_file) if os.path.isdir(_file): reader = SimpleDirectoryReader( @@ -228,7 +231,7 @@ class RagCLI(BaseModel): # Append the `--files` argument to the history file with open(f"{self.persist_dir}/{RAG_HISTORY_FILE_NAME}", "a") as f: - f.write(files + "\n") + f.write(str(files) + "\n") if create_llama: if shutil.which("npx") is None: @@ -337,9 +340,10 @@ class RagCLI(BaseModel): "-f", "--files", type=str, + nargs="+", help=( - "The name of the file or directory you want to ask a question about," - 'such as "file.pdf".' + "The name of the file(s) or directory you want to ask a question about," + 'such as "file.pdf". Supports globs like "*.py".' ), ) parser.add_argument( diff --git a/llama-index-cli/pyproject.toml b/llama-index-cli/pyproject.toml index 27a66f1e61c76e2da7eb8bd54fac476ec90680cd..812b8379f013c57a10b3f14ce3abc357e688085b 100644 --- a/llama-index-cli/pyproject.toml +++ b/llama-index-cli/pyproject.toml @@ -45,6 +45,7 @@ black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"} codespell = {extras = ["toml"], version = ">=v2.2.6"} ipython = "8.10.0" jupyter = "^1.0.0" +llama-index-vector-stores-chroma = "^0.4.1" mypy = "0.991" pre-commit = "3.2.0" pylint = "2.15.10" diff --git a/llama-index-cli/tests/test_rag.py b/llama-index-cli/tests/test_rag.py new file mode 100644 index 0000000000000000000000000000000000000000..64bd3c9f59b929e6fd1d5fd5f9ca6ab6f0956be2 --- /dev/null +++ b/llama-index-cli/tests/test_rag.py @@ -0,0 +1,16 @@ +import glob +from unittest import mock + + +from llama_index.cli import command_line +from llama_index.cli.rag import RagCLI + + +@mock.patch.object(RagCLI, "handle_cli", return_value="noop") +@mock.patch( + "sys.argv", + ["llamaindex-cli", "rag", "--files", *glob.glob("**/*.py", recursive=True)], +) +def test_handle_cli_files(mock_handle_cli) -> None: + command_line.main() + mock_handle_cli.assert_called_once()