Add JSON Lines support to JSONLoader (#6913)

**Description**: The JSON Lines format is used by some services such as OpenAI and HuggingFace. It's also a convenient alternative to CSV. This PR adds JSON Lines support to `JSONLoader` and also updates related tests. **Tag maintainer**: @rlancemartin, @eyurtsev. PS I was not able to build docs locally so didn't update related section.
2025-09-16 23:13:31 +00:00 · 2023-07-03 01:32:41 +06:00
parent 153b56d19b
commit 6d15854cda
5 changed files with 281 additions and 35 deletions
--- a/tests/unit_tests/document_loaders/test_json_loader.py
+++ b/tests/unit_tests/document_loaders/test_json_loader.py
@@ -1,3 +1,6 @@
+import io
+from typing import Any, Dict
+
 import pytest
 from pytest import raises
 from pytest_mock import MockerFixture
@@ -5,8 +8,9 @@ from pytest_mock import MockerFixture
 from langchain.docstore.document import Document
 from langchain.document_loaders.json_loader import JSONLoader

+pytestmark = pytest.mark.requires("jq")
+

-@pytest.mark.requires("jq")
 def test_load_valid_string_content(mocker: MockerFixture) -> None:
    file_path = "/workspaces/langchain/test.json"
    expected_docs = [
@@ -19,9 +23,12 @@ def test_load_valid_string_content(mocker: MockerFixture) -> None:
            metadata={"source": file_path, "seq_num": 2},
        ),
    ]
+
    mocker.patch("builtins.open", mocker.mock_open())
-    mock_csv_reader = mocker.patch("pathlib.Path.read_text")
-    mock_csv_reader.return_value = '[{"text": "value1"}, {"text": "value2"}]'
+    mocker.patch(
+        "pathlib.Path.read_text",
+        return_value='[{"text": "value1"}, {"text": "value2"}]',
+    )

    loader = JSONLoader(file_path=file_path, jq_schema=".[].text", text_content=True)
    result = loader.load()
@@ -29,7 +36,6 @@ def test_load_valid_string_content(mocker: MockerFixture) -> None:
    assert result == expected_docs


-@pytest.mark.requires("jq")
 def test_load_valid_dict_content(mocker: MockerFixture) -> None:
    file_path = "/workspaces/langchain/test.json"
    expected_docs = [
@@ -42,11 +48,14 @@ def test_load_valid_dict_content(mocker: MockerFixture) -> None:
            metadata={"source": file_path, "seq_num": 2},
        ),
    ]
+
    mocker.patch("builtins.open", mocker.mock_open())
-    mock_csv_reader = mocker.patch("pathlib.Path.read_text")
-    mock_csv_reader.return_value = """
-        [{"text": "value1"}, {"text": "value2"}]
-    """
+    mocker.patch(
+        "pathlib.Path.read_text",
+        return_value="""
+            [{"text": "value1"}, {"text": "value2"}]
+        """,
+    )

    loader = JSONLoader(file_path=file_path, jq_schema=".[]", text_content=False)
    result = loader.load()
@@ -54,7 +63,6 @@ def test_load_valid_dict_content(mocker: MockerFixture) -> None:
    assert result == expected_docs


-@pytest.mark.requires("jq")
 def test_load_valid_bool_content(mocker: MockerFixture) -> None:
    file_path = "/workspaces/langchain/test.json"
    expected_docs = [
@@ -67,13 +75,16 @@ def test_load_valid_bool_content(mocker: MockerFixture) -> None:
            metadata={"source": file_path, "seq_num": 2},
        ),
    ]
+
    mocker.patch("builtins.open", mocker.mock_open())
-    mock_csv_reader = mocker.patch("pathlib.Path.read_text")
-    mock_csv_reader.return_value = """
-        [
-            {"flag": false}, {"flag": true}
-        ]
-    """
+    mocker.patch(
+        "pathlib.Path.read_text",
+        return_value="""
+            [
+                {"flag": false}, {"flag": true}
+            ]
+        """,
+    )

    loader = JSONLoader(file_path=file_path, jq_schema=".[].flag", text_content=False)
    result = loader.load()
@@ -81,7 +92,6 @@ def test_load_valid_bool_content(mocker: MockerFixture) -> None:
    assert result == expected_docs


-@pytest.mark.requires("jq")
 def test_load_valid_numeric_content(mocker: MockerFixture) -> None:
    file_path = "/workspaces/langchain/test.json"
    expected_docs = [
@@ -94,13 +104,16 @@ def test_load_valid_numeric_content(mocker: MockerFixture) -> None:
            metadata={"source": file_path, "seq_num": 2},
        ),
    ]
+
    mocker.patch("builtins.open", mocker.mock_open())
-    mock_csv_reader = mocker.patch("pathlib.Path.read_text")
-    mock_csv_reader.return_value = """
-        [
-            {"num": 99}, {"num": 99.5}
-        ]
-    """
+    mocker.patch(
+        "pathlib.Path.read_text",
+        return_value="""
+            [
+                {"num": 99}, {"num": 99.5}
+            ]
+        """,
+    )

    loader = JSONLoader(file_path=file_path, jq_schema=".[].num", text_content=False)
    result = loader.load()
@@ -108,16 +121,152 @@ def test_load_valid_numeric_content(mocker: MockerFixture) -> None:
    assert result == expected_docs


-@pytest.mark.requires("jq")
 def test_load_invalid_test_content(mocker: MockerFixture) -> None:
    file_path = "/workspaces/langchain/test.json"
+
    mocker.patch("builtins.open", mocker.mock_open())
-    mock_csv_reader = mocker.patch("pathlib.Path.read_text")
-    mock_csv_reader.return_value = """
-        [{"text": "value1"}, {"text": "value2"}]
-    """
+    mocker.patch(
+        "pathlib.Path.read_text",
+        return_value="""
+            [{"text": "value1"}, {"text": "value2"}]
+        """,
+    )

    loader = JSONLoader(file_path=file_path, jq_schema=".[]", text_content=True)

    with raises(ValueError):
        loader.load()
+
+
+def test_load_jsonlines(mocker: MockerFixture) -> None:
+    file_path = "/workspaces/langchain/test.json"
+    expected_docs = [
+        Document(
+            page_content="value1",
+            metadata={"source": file_path, "seq_num": 1},
+        ),
+        Document(
+            page_content="value2",
+            metadata={"source": file_path, "seq_num": 2},
+        ),
+    ]
+
+    mocker.patch(
+        "pathlib.Path.open",
+        return_value=io.StringIO(
+            """
+            {"text": "value1"}
+            {"text": "value2"}
+            """
+        ),
+    )
+
+    loader = JSONLoader(
+        file_path=file_path, jq_schema=".", content_key="text", json_lines=True
+    )
+    result = loader.load()
+
+    assert result == expected_docs
+
+
+@pytest.mark.parametrize(
+    "params",
+    (
+        {"jq_schema": ".[].text"},
+        {"jq_schema": ".[]", "content_key": "text"},
+    ),
+)
+def test_load_jsonlines_list(params: Dict, mocker: MockerFixture) -> None:
+    file_path = "/workspaces/langchain/test.json"
+    expected_docs = [
+        Document(
+            page_content="value1",
+            metadata={"source": file_path, "seq_num": 1},
+        ),
+        Document(
+            page_content="value2",
+            metadata={"source": file_path, "seq_num": 2},
+        ),
+        Document(
+            page_content="value3",
+            metadata={"source": file_path, "seq_num": 3},
+        ),
+        Document(
+            page_content="value4",
+            metadata={"source": file_path, "seq_num": 4},
+        ),
+    ]
+
+    mocker.patch(
+        "pathlib.Path.open",
+        return_value=io.StringIO(
+            """
+            [{"text": "value1"}, {"text": "value2"}]
+            [{"text": "value3"}, {"text": "value4"}]
+            """
+        ),
+    )
+
+    loader = JSONLoader(file_path=file_path, json_lines=True, **params)
+    result = loader.load()
+
+    assert result == expected_docs
+
+
+def test_load_empty_jsonlines(mocker: MockerFixture) -> None:
+    mocker.patch("pathlib.Path.open", return_value=io.StringIO(""))
+
+    loader = JSONLoader(file_path="file_path", jq_schema=".[].text", json_lines=True)
+    result = loader.load()
+
+    assert result == []
+
+
+@pytest.mark.parametrize(
+    "patch_func,patch_func_value,kwargs",
+    (
+        # JSON content.
+        (
+            "pathlib.Path.read_text",
+            '[{"text": "value1"}, {"text": "value2"}]',
+            {"jq_schema": ".[]", "content_key": "text"},
+        ),
+        # JSON Lines content.
+        (
+            "pathlib.Path.open",
+            io.StringIO(
+                """
+                {"text": "value1"}
+                {"text": "value2"}
+                """
+            ),
+            {"jq_schema": ".", "content_key": "text", "json_lines": True},
+        ),
+    ),
+)
+def test_json_meta(
+    patch_func: str, patch_func_value: Any, kwargs: Dict, mocker: MockerFixture
+) -> None:
+    mocker.patch("builtins.open", mocker.mock_open())
+    mocker.patch(patch_func, return_value=patch_func_value)
+
+    file_path = "/workspaces/langchain/test.json"
+    expected_docs = [
+        Document(
+            page_content="value1",
+            metadata={"source": file_path, "seq_num": 1, "x": "value1-meta"},
+        ),
+        Document(
+            page_content="value2",
+            metadata={"source": file_path, "seq_num": 2, "x": "value2-meta"},
+        ),
+    ]
+
+    def metadata_func(record: Dict, metadata: Dict) -> Dict:
+        metadata["x"] = f"{record['text']}-meta"
+        return metadata
+
+    loader = JSONLoader(file_path=file_path, metadata_func=metadata_func, **kwargs)
+    result = loader.load()
+
+    assert result == expected_docs