feat (documents): add a source code loader based on AST manipulation (#6486)

#### Summary

A new approach to loading source code is implemented:

Each top-level function and class in the code is loaded into separate
documents. Then, an additional document is created with the top-level
code, but without the already loaded functions and classes.

This could improve the accuracy of QA chains over source code.

For instance, having this script:

```
class MyClass:
    def __init__(self, name):
        self.name = name

    def greet(self):
        print(f"Hello, {self.name}!")

def main():
    name = input("Enter your name: ")
    obj = MyClass(name)
    obj.greet()

if __name__ == '__main__':
    main()
```

The loader will create three documents with this content:

First document:
```
class MyClass:
    def __init__(self, name):
        self.name = name

    def greet(self):
        print(f"Hello, {self.name}!")
```

Second document:
```
def main():
    name = input("Enter your name: ")
    obj = MyClass(name)
    obj.greet()
```

Third document:
```
# Code for: class MyClass:

# Code for: def main():

if __name__ == '__main__':
    main()
```

A threshold parameter is added to control whether small scripts are
split in this way or not.

At this moment, only Python and JavaScript are supported. The
appropriate parser is determined by examining the file extension.

#### Tests

This PR adds:

- Unit tests
- Integration tests

#### Dependencies

Only one dependency was added as optional (needed for the JavaScript
parser).

#### Documentation

A notebook is added showing how the loader can be used.

#### Who can review?

@eyurtsev @hwchase17

---------

Co-authored-by: rlm <pexpresss31@gmail.com>

langchain/document_loaders/parsers/__init__.py

@@ -1,5 +1,6 @@
 from langchain.document_loaders.parsers.audio import OpenAIWhisperParser
 from langchain.document_loaders.parsers.html import BS4HTMLParser
+from langchain.document_loaders.parsers.language import LanguageParser
 from langchain.document_loaders.parsers.pdf import (
     PDFMinerParser,
     PDFPlumberParser,
@@ -10,6 +11,7 @@ from langchain.document_loaders.parsers.pdf import (
 
 __all__ = [
     "BS4HTMLParser",
+    "LanguageParser",
     "OpenAIWhisperParser",
     "PDFMinerParser",
     "PDFPlumberParser",

langchain/document_loaders/parsers/language/__init__.py

@@ -0,0 +1,3 @@
+from langchain.document_loaders.parsers.language.language_parser import LanguageParser
+
+__all__ = ["LanguageParser"]

langchain/document_loaders/parsers/language/code_segmenter.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+
+class CodeSegmenter(ABC):
+    def __init__(self, code: str):
+        self.code = code
+
+    def is_valid(self) -> bool:
+        return True
+
+    @abstractmethod
+    def simplify_code(self) -> str:
+        raise NotImplementedError  # pragma: no cover
+
+    @abstractmethod
+    def extract_functions_classes(self) -> List[str]:
+        raise NotImplementedError  # pragma: no cover

langchain/document_loaders/parsers/language/javascript.py

@@ -0,0 +1,65 @@
+from typing import Any, List
+
+from langchain.document_loaders.parsers.language.code_segmenter import CodeSegmenter
+
+
+class JavaScriptSegmenter(CodeSegmenter):
+    def __init__(self, code: str):
+        super().__init__(code)
+        self.source_lines = self.code.splitlines()
+
+        try:
+            import esprima  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "Could not import esprima Python package. "
+                "Please install it with `pip install esprima`."
+            )
+
+    def is_valid(self) -> bool:
+        import esprima
+
+        try:
+            esprima.parseScript(self.code)
+            return True
+        except esprima.Error:
+            return False
+
+    def _extract_code(self, node: Any) -> str:
+        start = node.loc.start.line - 1
+        end = node.loc.end.line
+        return "\n".join(self.source_lines[start:end])
+
+    def extract_functions_classes(self) -> List[str]:
+        import esprima
+
+        tree = esprima.parseScript(self.code, loc=True)
+        functions_classes = []
+
+        for node in tree.body:
+            if isinstance(
+                node,
+                (esprima.nodes.FunctionDeclaration, esprima.nodes.ClassDeclaration),
+            ):
+                functions_classes.append(self._extract_code(node))
+
+        return functions_classes
+
+    def simplify_code(self) -> str:
+        import esprima
+
+        tree = esprima.parseScript(self.code, loc=True)
+        simplified_lines = self.source_lines[:]
+
+        for node in tree.body:
+            if isinstance(
+                node,
+                (esprima.nodes.FunctionDeclaration, esprima.nodes.ClassDeclaration),
+            ):
+                start = node.loc.start.line - 1
+                simplified_lines[start] = f"// Code for: {simplified_lines[start]}"
+
+                for line_num in range(start + 1, node.loc.end.line):
+                    simplified_lines[line_num] = None  # type: ignore
+
+        return "\n".join(line for line in simplified_lines if line is not None)

langchain/document_loaders/parsers/language/language_parser.py

@@ -0,0 +1,143 @@
+from typing import Any, Dict, Iterator, Optional
+
+from langchain.docstore.document import Document
+from langchain.document_loaders.base import BaseBlobParser
+from langchain.document_loaders.blob_loaders import Blob
+from langchain.document_loaders.parsers.language.javascript import JavaScriptSegmenter
+from langchain.document_loaders.parsers.language.python import PythonSegmenter
+from langchain.text_splitter import Language
+
+LANGUAGE_EXTENSIONS: Dict[str, str] = {
+    "py": Language.PYTHON,
+    "js": Language.JS,
+}
+
+LANGUAGE_SEGMENTERS: Dict[str, Any] = {
+    Language.PYTHON: PythonSegmenter,
+    Language.JS: JavaScriptSegmenter,
+}
+
+
+class LanguageParser(BaseBlobParser):
+    """
+    Language parser that split code using the respective language syntax.
+
+    Each top-level function and class in the code is loaded into separate documents.
+    Furthermore, an extra document is generated, containing the remaining top-level code
+    that excludes the already segmented functions and classes.
+
+    This approach can potentially improve the accuracy of QA models over source code.
+
+    Currently, the supported languages for code parsing are Python and JavaScript.
+
+    The language used for parsing can be configured, along with the minimum number of
+    lines required to activate the splitting based on syntax.
+
+    Examples:
+
+        .. code-block:: python
+
+            from langchain.text_splitter.Language
+            from langchain.document_loaders.generic import GenericLoader
+            from langchain.document_loaders.parsers import LanguageParser
+
+            loader = GenericLoader.from_filesystem(
+                "./code",
+                glob="**/*",
+                suffixes=[".py", ".js"],
+                parser=LanguageParser()
+            )
+            docs = loader.load()
+
+        Example instantiations to manually select the language:
+
+        ... code-block:: python
+
+            from langchain.text_splitter import Language
+
+            loader = GenericLoader.from_filesystem(
+                "./code",
+                glob="**/*",
+                suffixes=[".py"],
+                parser=LanguageParser(language=Language.PYTHON)
+            )
+
+        Example instantiations to set number of lines threshold:
+
+        ... code-block:: python
+
+            loader = GenericLoader.from_filesystem(
+                "./code",
+                glob="**/*",
+                suffixes=[".py"],
+                parser=LanguageParser(parser_threshold=200)
+            )
+    """
+
+    def __init__(self, language: Optional[Language] = None, parser_threshold: int = 0):
+        """
+        Language parser that split code using the respective language syntax.
+
+        Args:
+            language: If None (default), it will try to infer language from source.
+            parser_threshold: Minimum lines needed to activate parsing (0 by default).
+        """
+        self.language = language
+        self.parser_threshold = parser_threshold
+
+    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
+        code = blob.as_string()
+
+        language = self.language or (
+            LANGUAGE_EXTENSIONS.get(blob.source.rsplit(".", 1)[-1])
+            if isinstance(blob.source, str)
+            else None
+        )
+
+        if language is None:
+            yield Document(
+                page_content=code,
+                metadata={
+                    "source": blob.source,
+                },
+            )
+            return
+
+        if self.parser_threshold >= len(code.splitlines()):
+            yield Document(
+                page_content=code,
+                metadata={
+                    "source": blob.source,
+                    "language": language,
+                },
+            )
+            return
+
+        self.Segmenter = LANGUAGE_SEGMENTERS[language]
+        segmenter = self.Segmenter(blob.as_string())
+        if not segmenter.is_valid():
+            yield Document(
+                page_content=code,
+                metadata={
+                    "source": blob.source,
+                },
+            )
+            return
+
+        for functions_classes in segmenter.extract_functions_classes():
+            yield Document(
+                page_content=functions_classes,
+                metadata={
+                    "source": blob.source,
+                    "content_type": "functions_classes",
+                    "language": language,
+                },
+            )
+        yield Document(
+            page_content=segmenter.simplify_code(),
+            metadata={
+                "source": blob.source,
+                "content_type": "simplified_code",
+                "language": language,
+            },
+        )

langchain/document_loaders/parsers/language/python.py

@@ -0,0 +1,47 @@
+import ast
+from typing import Any, List
+
+from langchain.document_loaders.parsers.language.code_segmenter import CodeSegmenter
+
+
+class PythonSegmenter(CodeSegmenter):
+    def __init__(self, code: str):
+        super().__init__(code)
+        self.source_lines = self.code.splitlines()
+
+    def is_valid(self) -> bool:
+        try:
+            ast.parse(self.code)
+            return True
+        except SyntaxError:
+            return False
+
+    def _extract_code(self, node: Any) -> str:
+        start = node.lineno - 1
+        end = node.end_lineno
+        return "\n".join(self.source_lines[start:end])
+
+    def extract_functions_classes(self) -> List[str]:
+        tree = ast.parse(self.code)
+        functions_classes = []
+
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                functions_classes.append(self._extract_code(node))
+
+        return functions_classes
+
+    def simplify_code(self) -> str:
+        tree = ast.parse(self.code)
+        simplified_lines = self.source_lines[:]
+
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                start = node.lineno - 1
+                simplified_lines[start] = f"# Code for: {simplified_lines[start]}"
+
+                assert isinstance(node.end_lineno, int)
+                for line_num in range(start + 1, node.end_lineno):
+                    simplified_lines[line_num] = None  # type: ignore
+
+        return "\n".join(line for line in simplified_lines if line is not None)