Fix chunking for .ts and .tsx files (#73)

pyproject.toml

@@ -56,6 +56,7 @@ dependencies = [
     "tokenizers==0.19.1",
     "transformers==4.44.2",
     "tree-sitter==0.22.3",
+    "tree-sitter-typescript==0.21.2",
     "tree-sitter-language-pack==0.2.0",
     "voyageai==0.2.3",
     "setuptools"  # Added from the setup.py install_requires

sage/chunker.py

@@ -170,10 +170,7 @@ class CodeFileChunker(Chunker):
     def is_code_file(filename: str) -> bool:
         """Checks whether pygment & tree_sitter can parse the file as code."""
         language = CodeFileChunker._get_language_from_filename(filename)
-        # tree-sitter-language-pack crashes on TypeScript files. We'll wait for a bit to see if the issue gets
-        # resolved, otherwise we'll have to clone and fix the library.
-        # See https://github.com/Goldziher/tree-sitter-language-pack/issues/5
-        return language and language not in ["text only", "None", "typescript", "tsx"]
+        return language and language not in ["text only", "None"]
 
     @staticmethod
     def parse_tree(filename: str, content: str) -> List[str]:
@@ -184,12 +181,6 @@ class CodeFileChunker(Chunker):
             logging.debug("%s doesn't seem to be a code file.", filename)
             return None
 
-        if language in ["typescript", "tsx"]:
-            # tree-sitter-language-pack crashes on TypeScript files. We'll wait for a bit to see if the issue gets
-            # resolved, otherwise we'll have to clone and fix the library.
-            # See https://github.com/Goldziher/tree-sitter-language-pack/issues/5
-            return None
-
         try:
             parser = get_parser(language)
         except LookupError:

tests/test_chunker.py

@@ -40,9 +40,11 @@ def test_code_chunker_happy_path():
     assert len(chunks) >= 1
 
 
-@mark.parametrize("filename", [param("assets/sample-script.ts"), param("assets/sample-script.tsx")])
-def test_code_chunker_typescript(filename):
-    """Tests CodeFileChunker on .ts and .tsx files (tree_sitter_language_pack doesn't work out of the box)."""
+@mark.parametrize(
+    "filename", [param("assets/sample-script.ts"), param("assets/sample-script.tsx")]
+)
+def test_code_chunker_typescript_happy_path(filename):
+    """Tests the happy path for the CodeFileChunker on .ts and .tsx files."""
     file_path = os.path.join(os.path.dirname(__file__), filename)
     with open(file_path, "r") as file:
         content = file.read()
@@ -50,13 +52,7 @@ def test_code_chunker_typescript(filename):
 
     chunker = sage.chunker.CodeFileChunker(max_tokens=100)
     chunks = chunker.chunk(content, metadata)
-    # There's a bug in the tree-sitter-language-pack library for TypeScript. Before it gets fixed, we expect this to
-    # return an empty list (instead of crashing).
-    assert len(chunks) == 0
 
-    # However, the UniversalFileChunker should fallback onto a regular text chunker, and return some chunks.
-    chunker = sage.chunker.UniversalFileChunker(max_tokens=100)
-    chunks = chunker.chunk(content, metadata)
     assert len(chunks) >= 1