Add file filter param to Git loader (#2904)

Allows users to specify what files should be loaded instead of
indiscriminately loading the entire repo.

extends #2851 

NOTE: for reviewers, `hide whitespace` option recommended since I
changed the indentation of an if-block to use `continue` instead so it
looks less like a Christmas tree :)
This commit is contained in:
Kwuang Tang 2023-04-14 13:45:54 -04:00 committed by GitHub
parent 7e525a3b91
commit a508afa91c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 65 additions and 47 deletions

@ -0,0 +1 @@
Subproject commit 7e525a3b91ddf2cc0d9b3add5928460449c1d600

View File

@ -18,7 +18,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 21, "execution_count": 5,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -32,7 +32,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 22, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -41,7 +41,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 23, "execution_count": 6,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -50,7 +50,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 24, "execution_count": 8,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -59,27 +59,16 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 25, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [],
{
"data": {
"text/plain": [
"1040"
]
},
"execution_count": 25,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [ "source": [
"len(data)" "len(data)"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 26, "execution_count": 9,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [
{ {
@ -103,7 +92,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 27, "execution_count": 5,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -112,7 +101,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 28, "execution_count": 2,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -125,7 +114,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 29, "execution_count": 10,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -134,16 +123,16 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 30, "execution_count": 11,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [
{ {
"data": { "data": {
"text/plain": [ "text/plain": [
"1040" "1074"
] ]
}, },
"execution_count": 30, "execution_count": 11,
"metadata": {}, "metadata": {},
"output_type": "execute_result" "output_type": "execute_result"
} }
@ -152,6 +141,25 @@
"len(data)" "len(data)"
] ]
}, },
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Filtering files to load"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"from langchain.document_loaders import GitLoader\n",
"\n",
"# eg. loading only python files\n",
"loader = GitLoader(repo_path=\"./example_data/test_repo1/\", file_filter=lambda file_path: file_path.endswith(\".py\"))"
]
},
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": null,
@ -176,7 +184,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.9.1" "version": "3.11.3"
} }
}, },
"nbformat": 4, "nbformat": 4,

View File

@ -1,5 +1,5 @@
import os import os
from typing import List, Optional from typing import Callable, List, Optional
from langchain.docstore.document import Document from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader from langchain.document_loaders.base import BaseLoader
@ -21,10 +21,12 @@ class GitLoader(BaseLoader):
repo_path: str, repo_path: str,
clone_url: Optional[str] = None, clone_url: Optional[str] = None,
branch: Optional[str] = "main", branch: Optional[str] = "main",
file_filter: Optional[Callable[[str], bool]] = None,
): ):
self.repo_path = repo_path self.repo_path = repo_path
self.clone_url = clone_url self.clone_url = clone_url
self.branch = branch self.branch = branch
self.file_filter = file_filter
def load(self) -> List[Document]: def load(self) -> List[Document]:
try: try:
@ -47,8 +49,15 @@ class GitLoader(BaseLoader):
docs: List[Document] = [] docs: List[Document] = []
for item in repo.tree().traverse(): for item in repo.tree().traverse():
if isinstance(item, Blob): if not isinstance(item, Blob):
continue
file_path = os.path.join(self.repo_path, item.path) file_path = os.path.join(self.repo_path, item.path)
# uses filter to skip files
if self.file_filter and not self.file_filter(file_path):
continue
rel_file_path = os.path.relpath(file_path, self.repo_path) rel_file_path = os.path.relpath(file_path, self.repo_path)
try: try:
with open(file_path, "rb") as f: with open(file_path, "rb") as f: