From 3f0a651c8e57b101ebb1e7a20e6cf2d59e1dc9c1 Mon Sep 17 00:00:00 2001
From: mohammadrezapourreza <m1378.prz@gmail.com>
Date: Wed, 13 Sep 2023 19:04:31 -0400
Subject: [PATCH] DH-4654/adding google api exceptions to the agent

---
 dataherald/sql_generator/dataherald_sqlagent.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py
index 2e775903..d13e62ab 100644
--- a/dataherald/sql_generator/dataherald_sqlagent.py
+++ b/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}"