diff --git a/libs/partners/upstage/langchain_upstage/layout_analysis_parsers.py b/libs/partners/upstage/langchain_upstage/layout_analysis_parsers.py index 5a4056dfe61..17b571c5132 100644 --- a/libs/partners/upstage/langchain_upstage/layout_analysis_parsers.py +++ b/libs/partners/upstage/langchain_upstage/layout_analysis_parsers.py @@ -181,19 +181,22 @@ class UpstageLayoutAnalysisParser(BaseBlobParser): result = response.json().get("elements", []) + elements = [ + element for element in result if element["category"] not in self.exclude + ] + + return elements + except requests.RequestException as req_err: # Handle any request-related exceptions print(f"Request Exception: {req_err}") + raise ValueError(f"Failed to send request: {req_err}") except json.JSONDecodeError as json_err: # Handle JSON decode errors print(f"JSON Decode Error: {json_err}") raise ValueError(f"Failed to decode JSON response: {json_err}") - elements = [ - element for element in result if element["category"] not in self.exclude - ] - - return elements + return [] def _split_and_request( self, diff --git a/libs/partners/upstage/tests/unit_tests/test_layout_analysis.py b/libs/partners/upstage/tests/unit_tests/test_layout_analysis.py index a74b914fb07..86c5cbb0d90 100644 --- a/libs/partners/upstage/tests/unit_tests/test_layout_analysis.py +++ b/libs/partners/upstage/tests/unit_tests/test_layout_analysis.py @@ -1,7 +1,11 @@ +import json from pathlib import Path from typing import Any, Dict, get_args +from unittest import TestCase from unittest.mock import MagicMock, Mock, patch +import requests + from langchain_upstage import UpstageLayoutAnalysisLoader from langchain_upstage.layout_analysis import OutputType, SplitType @@ -205,3 +209,45 @@ def test_page_split_html_output(mock_post: Mock) -> None: assert document.metadata["page"] == MOCK_RESPONSE_JSON["elements"][i]["page"] assert document.metadata["type"] == "html" assert document.metadata["split"] == "page" + + +@patch("requests.post") +def test_request_exception(mock_post: Mock) -> None: + mock_post.side_effect = requests.RequestException("Mocked request exception") + + loader = UpstageLayoutAnalysisLoader( + file_path=EXAMPLE_PDF_PATH, + output_type="html", + split="page", + api_key="valid_api_key", + exclude=[], + ) + + with TestCase.assertRaises(TestCase(), ValueError) as context: + loader.load() + + assert "Failed to send request: Mocked request exception" == str(context.exception) + + +@patch("requests.post") +def test_json_decode_error(mock_post: Mock) -> None: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + mock_post.return_value = mock_response + + loader = UpstageLayoutAnalysisLoader( + file_path=EXAMPLE_PDF_PATH, + output_type="html", + split="page", + api_key="valid_api_key", + exclude=[], + ) + + with TestCase.assertRaises(TestCase(), ValueError) as context: + loader.load() + + assert ( + "Failed to decode JSON response: Expecting value: line 1 column 1 (char 0)" + == str(context.exception) + )