diff --git a/api/config.py b/api/config.py index d6f9232a7..ab9877271 100644 --- a/api/config.py +++ b/api/config.py @@ -417,6 +417,105 @@ def get_model_config(provider="google", model=None): return result +def _should_process_file( + file_path: Path, + use_inclusion: bool, + included_dirs: list[str], + included_files: list[str], + excluded_dirs: list[str], + excluded_files: list[str], +) -> bool: + """Decide if a file passes the include/exclude rules (moved from rag.pipeline + so the tree listing and the RAG indexer share one implementation).""" + if isinstance(file_path, str): + file_path = Path(file_path) + file_path_parts = file_path.resolve().parts + file_name = file_path_parts[-1] + + if use_inclusion: + is_included = False + if included_dirs: + for included in included_dirs: + clean_included = included.removeprefix("./").rstrip("/") + if clean_included in file_path_parts: + is_included = True + break + if not is_included and included_files: + for included_file in included_files: + if file_name == included_file or file_name.endswith(included_file): + is_included = True + break + if not included_dirs and not included_files: + is_included = True + return is_included + + is_excluded = False + if excluded_dirs: + for excluded in excluded_dirs: + clean_excluded = excluded.removeprefix("./").rstrip("/") + if clean_excluded in file_path_parts: + is_excluded = True + break + if not is_excluded and excluded_files: + for excluded_file in excluded_files: + if file_name == excluded_file: + is_excluded = True + break + return not is_excluded + + +def iterate_files( + root_dir: str, + excluded_dirs: list[str] | None = None, + excluded_files: list[str] | None = None, + included_dirs: list[str] | None = None, + included_files: list[str] | None = None, +) -> list[str]: + """Walk ``root_dir`` and return repo-relative paths of the files worth + processing, using the SAME rules the RAG indexer uses so the wiki-structure + file tree matches what actually gets indexed: + + * restrict to the configured code/doc extensions; + * exclusion mode: config ``file_filters`` excluded_dirs/files UNION the + request-provided excluded_dirs/files; + * inclusion mode (when included_dirs/files are given): only those. + """ + use_inclusion = bool(included_dirs or included_files) + if use_inclusion: + inc_dirs = list(set(included_dirs or [])) + inc_files = list(set(included_files or [])) + exc_dirs: list[str] = [] + exc_files: list[str] = [] + else: + file_filters = configs.get("file_filters", {}) + exc_dir_set = set(file_filters.get("excluded_dirs", [])) + exc_file_set = set(file_filters.get("excluded_files", [])) + if excluded_dirs: + exc_dir_set.update(excluded_dirs) + if excluded_files: + exc_file_set.update(excluded_files) + exc_dirs = list(exc_dir_set) + exc_files = list(exc_file_set) + inc_dirs = [] + inc_files = [] + + extensions = tuple( + configs.get("code_extensions", []) + configs.get("doc_extensions", []) + ) + + results: list[str] = [] + for p in Path(root_dir).rglob("*"): + if not p.is_file(): + continue + if extensions and p.suffix.lower() not in extensions: + continue + if _should_process_file( + p, use_inclusion, inc_dirs, inc_files, exc_dirs, exc_files + ): + results.append(os.path.relpath(p, root_dir).replace(os.sep, "/")) + return results + + def get_embedder( is_local_ollama: bool = False, use_google_embedder: bool = False, diff --git a/api/rag/pipeline.py b/api/rag/pipeline.py index ab9a34e5c..83b63783c 100644 --- a/api/rag/pipeline.py +++ b/api/rag/pipeline.py @@ -10,6 +10,7 @@ from api.config import ( configs, get_embedder, + iterate_files, ) from api.logger import get_logger from api.repository import Repo @@ -69,85 +70,6 @@ def count_tokens( return len(text) // 4 -def _should_process_file( - file_path: Path, - use_inclusion: bool, - included_dirs: list[str] | None, - included_files: list[str] | None, - excluded_dirs: list[str] | None, - excluded_files: list[str] | None, -) -> bool: - """ - Determine if a file should be processed based on inclusion/exclusion rules. - - Args: - file_path (str): The file path to check - use_inclusion (bool): Whether to use inclusion mode - included_dirs (List[str]): List of directories to include - included_files (List[str]): List of files to include - excluded_dirs (List[str]): List of directories to exclude - excluded_files (List[str]): List of files to exclude - - Returns: - bool: True if the file should be processed, False otherwise - """ - if isinstance(file_path, str): - file_path = Path(file_path) - file_path_parts = file_path.resolve().parts - file_name = file_path_parts[-1] - - if use_inclusion: - # Inclusion mode: file must be in included directories or match included files - is_included = False - - # Check if file is in an included directory - if included_dirs: - for included in included_dirs: - clean_included = included.removeprefix("./").rstrip("/") - if clean_included in file_path_parts: - is_included = True - break - - # Check if file matches included file patterns - if not is_included and included_files: - for included_file in included_files: - if file_name == included_file or file_name.endswith(included_file): - is_included = True - break - - # If no inclusion rules are specified for a category, allow all files from that category - if not included_dirs and not included_files: - is_included = True - elif not included_dirs and included_files: - # Only file patterns specified, allow all directories - pass # is_included is already set based on file patterns - elif included_dirs and not included_files: - # Only directory patterns specified, allow all files in included directories - pass # is_included is already set based on directory patterns - - return is_included - else: - # Exclusion mode: file must not be in excluded directories or match excluded files - is_excluded = False - - # Check if file is in an excluded directory - if excluded_dirs: - for excluded in excluded_dirs: - clean_excluded = excluded.removeprefix("./").rstrip("/") - if clean_excluded in file_path_parts: - is_excluded = True - break - - # Check if file matches excluded file patterns - if not is_excluded and excluded_files: - for excluded_file in excluded_files: - if file_name == excluded_file: - is_excluded = True - break - - return not is_excluded - - def read_all_documents( path: str, embedder_type: str = None, @@ -182,68 +104,22 @@ def read_all_documents( if embedder_type is None and is_ollama_embedder is not None: embedder_type = "ollama" if is_ollama_embedder else None documents = [] - # File extensions to look for, prioritizing code files code_extensions = configs.get("code_extensions", []) - doc_extensions = configs.get("doc_extensions", []) - - # Determine filtering mode: inclusion or exclusion - use_inclusion_mode = bool(included_dirs or included_files) - - if use_inclusion_mode: - # Inclusion mode: only process specified directories and files - included_dirs = list(set(included_dirs)) if included_dirs else list() - included_files = list(set(included_files)) if included_files else list() - - logger.info("Using inclusion mode") - logger.info(f"Included directories: {included_dirs}") - logger.info(f"Included files: {included_files}") - - # Convert to lists for processing - excluded_dirs = [] - excluded_files = [] - else: - # Exclusion mode: use default exclusions plus any additional ones - file_filters = configs.get("file_filters", {}) - final_excluded_dirs: set[str] = set(file_filters.get("excluded_dirs", [])) - final_excluded_files: set[str] = set(file_filters.get("excluded_files", [])) - - # Add any explicitly provided excluded directories and files - if excluded_dirs is not None: - final_excluded_dirs.update(excluded_dirs) - - if excluded_files is not None: - final_excluded_files.update(excluded_files) - - # Convert back to lists for compatibility - excluded_dirs = list(final_excluded_dirs) - excluded_files = list(final_excluded_files) - included_dirs = [] - included_files = [] - - logger.info("Using exclusion mode") - logger.info(f"Excluded directories: {excluded_dirs}") - logger.info(f"Excluded files: {excluded_files}") logger.info(f"Reading documents from {path}") - for file_path in filter( - lambda p: _should_process_file( - p, - use_inclusion=use_inclusion_mode, - included_dirs=included_dirs, - included_files=included_files, - excluded_dirs=excluded_dirs, - excluded_files=excluded_files, - ), - filter( - lambda p: p.suffix.lower() in code_extensions + doc_extensions, - Path(path).rglob(pattern="**/*"), - ), + # Single source of truth for which files to process (see config.iterate_files). + for relative_path in iterate_files( + path, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + included_dirs=included_dirs, + included_files=included_files, ): + file_path = Path(path) / relative_path try: with open(file_path, "r", encoding="utf-8") as f: content = f.read() - relative_path = os.path.relpath(file_path, path) # Check token count token_count = count_tokens(content, embedder_type) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 000000000..a2197fd33 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,82 @@ +import pytest + +from api.config import configs, iterate_files + + +@pytest.fixture +def patched_config(monkeypatch): + monkeypatch.setitem(configs, "code_extensions", [".py"]) + monkeypatch.setitem(configs, "doc_extensions", [".md"]) + monkeypatch.setitem( + configs, + name="file_filters", + value={ + "excluded_dirs": [ + "./.venv/", + "./venv/", + ], + "excluded_files": [ + "yarn.lock", + ".env", + ] + } + ) + + +def make_repo(root): + (root / "README.md").write_text("") + (root / "CHANGELOG.md").write_text("") + (root / "yarn.lock").write_text("") # excluded + + folder = root / "folder" + folder.mkdir(exist_ok=True) + (folder / ".lock").write_text("") # excluded + (folder / "code.py").write_text("") + + ex_folder = root / ".venv" + ex_folder.mkdir(exist_ok=True) + (ex_folder / "file.txt").write_text("") + (ex_folder / ".gitignore").write_text("") + + + +def test_iterate_files_default_exclusive_mode(patched_config, tmp_path): + make_repo(tmp_path) + + files = set(iterate_files(root_dir=str(tmp_path))) + assert files == { + "README.md", + "CHANGELOG.md", + "folder/code.py", + } + +@pytest.mark.parametrize( + "included_dirs", + [ + ["folder"], + ["./folder"], + ] +) +def test_iterate_files_included_dirs(patched_config, tmp_path, included_dirs): + make_repo(tmp_path) + files = set(iterate_files(root_dir=str(tmp_path), included_dirs=included_dirs)) + assert files == {"folder/code.py"} + + +def test_iterate_files_included_files(patched_config, tmp_path): + make_repo(tmp_path) + files = set(iterate_files(root_dir=str(tmp_path), included_files=["README.md"])) + assert files == {"README.md"} + + +@pytest.mark.parametrize( + "excluded_dirs", + [ + ["folder"], + ["./folder"], + ] +) +def test_iterate_files_excluded_dirs(patched_config, tmp_path, excluded_dirs): + make_repo(tmp_path) + files = set(iterate_files(root_dir=str(tmp_path), excluded_dirs=excluded_dirs)) + assert files == {"README.md", "CHANGELOG.md"}