feat: Raise an exception when get error HTTP code (#646)

pandasai/exceptions.py

@@ -68,6 +68,24 @@ class MissingModelError(Exception):
     """
 
 
+class LLMResponseHTTPError(Exception):
+    """
+    Raised when a remote LLM service responses with error HTTP code.
+
+    Args:
+        Exception (Exception): LLMResponseHTTPError
+    """
+
+    def __init__(self, status_code, error_msg=None):
+        self.status_code = status_code
+        self.error_msg = error_msg
+
+        super().__init__(
+            f"The remote server has responded with an error HTTP "
+            f"code: {status_code}; {error_msg or ''}"
+        )
+
+
 class BadImportError(Exception):
     """
     Raised when a library not in the whitelist is imported.

pandasai/llm/base.py

@@ -27,6 +27,7 @@ class CustomLLM(BaseOpenAI):
     APIKeyNotFoundError,
     MethodNotImplementedError,
     NoCodeFoundError,
+    LLMResponseHTTPError,
 )
 from ..helpers.openai_info import openai_callback_var
 from ..prompts.base import AbstractPrompt
@@ -324,6 +325,10 @@ def query(self, payload) -> str:
             str: Value of the field "generated_text" in response JSON
                 given by the remote server.
 
+        Raises:
+            LLMResponseHTTPError: If api-inference.huggingface.co responses
+                with any error HTTP code (>= 400).
+
         """
 
         headers = {"Authorization": f"Bearer {self.api_token}"}
@@ -332,6 +337,16 @@ def query(self, payload) -> str:
             self._api_url, headers=headers, json=payload, timeout=60
         )
 
+        if response.status_code >= 400:
+            try:
+                error_msg = response.json().get("error")
+            except (requests.exceptions.JSONDecodeError, TypeError):
+                error_msg = None
+
+            raise LLMResponseHTTPError(
+                status_code=response.status_code, error_msg=error_msg
+            )
+
         return response.json()[0]["generated_text"]
 
     def call(self, instruction: AbstractPrompt, suffix: str = "") -> str:

tests/llms/test_base_hf.py

@@ -3,6 +3,7 @@
 import pytest
 import requests
 
+from pandasai.exceptions import LLMResponseHTTPError
 from pandasai.llm.base import HuggingFaceLLM
 from pandasai.prompts import AbstractPrompt
 
@@ -14,6 +15,10 @@ class TestBaseHfLLM:
     def api_response(self):
         return [{"generated_text": "Some text"}]
 
+    @pytest.fixture
+    def api_response_401(self):
+        return {"error": "Authorization header is correct, but the token seems invalid"}
+
     @pytest.fixture
     def prompt(self):
         class MockAbstractPrompt(AbstractPrompt):
@@ -32,6 +37,7 @@ def test_api_url(self):
 
     def test_query(self, mocker, api_response):
         response_mock = mocker.Mock()
+        response_mock.status_code = 200
         response_mock.json.return_value = api_response
         mocker.patch("requests.post", return_value=response_mock)
 
@@ -51,6 +57,27 @@ def test_query(self, mocker, api_response):
         # Check that the result is correct
         assert result == api_response[0]["generated_text"]
 
+    def test_query_http_error_401(self, mocker, api_response_401):
+        response_mock = mocker.Mock()
+        response_mock.status_code = 401
+        response_mock.json.return_value = api_response_401
+        mocker.patch("requests.post", return_value=response_mock)
+
+        llm = HuggingFaceLLM(api_token="test_token")
+        payload = {"inputs": "Some input text"}
+
+        with pytest.raises(LLMResponseHTTPError) as exc:
+            llm.query(payload)
+
+        assert api_response_401.get("error") in str(exc.value)
+
+        requests.post.assert_called_once_with(
+            llm._api_url,
+            headers={"Authorization": "Bearer test_token"},
+            json=payload,
+            timeout=60,
+        )
+
     def test_call(self, mocker, prompt):
         huggingface = HuggingFaceLLM(api_token="test_token")