DH-4654/adding google api exceptions to the agent

dataherald/sql_generator/dataherald_sqlagent.py

@@ -8,6 +8,7 @@
 import numpy as np
 import openai
 import pandas as pd
+from google.api_core.exceptions import GoogleAPIError
 from langchain.agents.agent import AgentExecutor
 from langchain.agents.agent_toolkits.base import BaseToolkit
 from langchain.agents.mrkl.base import ZeroShotAgent
@@ -87,7 +88,7 @@
 
 
 def catch_exceptions(max_handling_list):  # noqa: C901
-    def decorator(fn: Callable[[str], str]) -> Callable[[str], str]:
+    def decorator(fn: Callable[[str], str]) -> Callable[[str], str]:  # noqa: C901
         @wraps(fn)
         def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: PLR0911
             nonlocal max_handling_list
@@ -116,6 +117,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: PLR0911
                     return f"OpenAI API service unavailable: {e}"
                 except openai.error.InvalidRequestError as e:
                     return f"OpenAI API request was invalid: {e}"
+                except GoogleAPIError as e:
+                    return f"Google API returned an error: {e}"
                 except SQLAlchemyError as e:
                     return f"Error: {e}"