fix: Error handling improvement for gemini (#84)

* Error handling fix for gemini, bumped SDK to 0.41.1

* Update src/unstract/sdk/adapters/llm/vertex_ai/src/vertex_ai.py

Co-authored-by: Deepak K <[email protected]>
Signed-off-by: Chandrasekharan M <[email protected]>

* Minor fixes on gemini LLM error handling

---------

Signed-off-by: Chandrasekharan M <[email protected]>
Co-authored-by: Deepak K <[email protected]>

Author: chandrasekharan-zipstack

src/unstract/sdk/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.41.0"
+__version__ = "0.41.1"
 
 
 def get_sdk_version():

src/unstract/sdk/adapter.py

@@ -58,6 +58,7 @@ def get_adapter_configuration(
         if response.status_code == 200:
             adapter_data: dict[str, Any] = response.json()
 
+            # TODO: Print config after redacting sensitive information
             self.tool.stream_log(
                 "Successfully retrieved adapter config "
                 f"for adapter: {adapter_instance_id}"
@@ -104,18 +105,16 @@ def get_adapter_config(
             Any: engine
         """
         # Check if the adapter ID matches any public adapter keys
-        if SdkHelper.is_public_adapter(
-            adapter_id=adapter_instance_id
-        ):
-            adapter_metadata_config = tool.get_env_or_die(
-                adapter_instance_id
-            )
+        if SdkHelper.is_public_adapter(adapter_id=adapter_instance_id):
+            adapter_metadata_config = tool.get_env_or_die(adapter_instance_id)
             adapter_metadata = json.loads(adapter_metadata_config)
             return adapter_metadata
         platform_host = tool.get_env_or_die(ToolEnv.PLATFORM_HOST)
         platform_port = tool.get_env_or_die(ToolEnv.PLATFORM_PORT)
 
-        tool.stream_log("Connecting to DB and getting table metadata")
+        tool.stream_log(
+            f"Connecting to DB and getting table metadata for {adapter_instance_id}"
+        )
         tool_adapter = ToolAdapter(
             tool=tool,
             platform_host=platform_host,

src/unstract/sdk/adapters/exceptions.py

@@ -1,14 +1,8 @@
-from unstract.sdk.adapters.constants import Common
+from unstract.sdk.exceptions import SdkError
 
 
-class AdapterError(Exception):
-    def __init__(self, message: str = Common.DEFAULT_ERR_MESSAGE):
-        super().__init__(message)
-        # Make it user friendly wherever possible
-        self.message = message
-
-    def __str__(self) -> str:
-        return self.message
+class AdapterError(SdkError):
+    pass
 
 
 class LLMError(AdapterError):

src/unstract/sdk/adapters/llm/exceptions.py

@@ -0,0 +1,25 @@
+from openai import APIError as OpenAIAPIError
+from vertexai.generative_models import ResponseValidationError
+
+from unstract.sdk.adapters.exceptions import LLMError
+from unstract.sdk.adapters.llm.open_ai.src.open_ai import OpenAILLM
+from unstract.sdk.adapters.llm.vertex_ai.src.vertex_ai import VertexAILLM
+
+
+def parse_llm_err(e: Exception) -> LLMError:
+    """Parses the exception from LLM provider.
+
+    Helps parse the LLM error and wraps it with our
+    custom exception object to contain a user friendly message.
+
+    Args:
+        e (Exception): Error from LLM provider
+
+    Returns:
+        LLMError: Unstract's LLMError object
+    """
+    if isinstance(e, ResponseValidationError):
+        return VertexAILLM.parse_llm_err(e)
+    elif isinstance(e, OpenAIAPIError):
+        return OpenAILLM.parse_llm_err(e)
+    return LLMError(str(e))

src/unstract/sdk/adapters/llm/llm_adapter.py

@@ -6,6 +6,7 @@
 
 from unstract.sdk.adapters.base import Adapter
 from unstract.sdk.adapters.enums import AdapterTypes
+from unstract.sdk.adapters.exceptions import LLMError
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +45,20 @@ def get_json_schema() -> str:
     def get_adapter_type() -> AdapterTypes:
         return AdapterTypes.LLM
 
+    @staticmethod
+    def parse_llm_err(e: Exception) -> LLMError:
+        """Parse the error from an LLM provider.
+
+        Helps parse errors from a provider and wraps with custom exception.
+
+        Args:
+            e (Exception): Exception from LLM provider
+
+        Returns:
+            LLMError: Error to be sent to the user
+        """
+        return LLMError(str(e))
+
     def get_llm_instance(self) -> LLM:
         """Instantiate the llama index LLM class.
 

src/unstract/sdk/adapters/llm/open_ai/src/open_ai.py

@@ -3,11 +3,14 @@
 
 from llama_index.core.llms import LLM
 from llama_index.llms.openai import OpenAI
+from openai import APIError as OpenAIAPIError
+from openai import RateLimitError as OpenAIRateLimitError
 
 from unstract.sdk.adapters.exceptions import AdapterError
 from unstract.sdk.adapters.llm.constants import LLMKeys
 from unstract.sdk.adapters.llm.helper import LLMHelper
 from unstract.sdk.adapters.llm.llm_adapter import LLMAdapter
+from unstract.sdk.exceptions import LLMError, RateLimitError
 
 
 class Constants:
@@ -76,3 +79,24 @@ def test_connection(self) -> bool:
         llm = self.get_llm_instance()
         test_result: bool = LLMHelper.test_llm_instance(llm=llm)
         return test_result
+
+    @staticmethod
+    def parse_llm_err(e: OpenAIAPIError) -> LLMError:
+        """Parse the error from Open AI.
+
+        Helps parse errors from Open AI and wraps with custom exception.
+
+        Args:
+            e (OpenAIAPIError): Exception from Open AI
+
+        Returns:
+            LLMError: Error to be sent to the user
+        """
+        msg = "OpenAI error: "
+        if hasattr(e, "body") and "message" in e.body:
+            msg += e.body["message"]
+        else:
+            msg += e.message
+        if isinstance(e, OpenAIRateLimitError):
+            return RateLimitError(msg)
+        return LLMError(msg)

src/unstract/sdk/adapters/llm/vertex_ai/src/vertex_ai.py

@@ -7,6 +7,7 @@
 from google.oauth2.service_account import Credentials
 from llama_index.core.llms import LLM
 from llama_index.llms.vertex import Vertex
+from vertexai.generative_models import Candidate, FinishReason, ResponseValidationError
 from vertexai.generative_models._generative_models import (
     HarmBlockThreshold,
     HarmCategory,
@@ -191,3 +192,62 @@ def test_connection(self) -> bool:
         except Exception as e:
             raise LLMError(f"Error while testing connection for VertexAI: {str(e)}")
         return test_result
+
+    @staticmethod
+    def parse_llm_err(e: ResponseValidationError) -> LLMError:
+        """Parse the error from Vertex AI.
+
+        Helps parse and raise errors from Vertex AI.
+        https://ai.google.dev/api/generate-content#generatecontentresponse
+
+        Args:
+            e (ResponseValidationError): Exception from Vertex AI
+
+        Returns:
+            LLMError: Error to be sent to the user
+        """
+        assert len(e.responses) == 1, (
+            "Expected e.responses to contain a single element "
+            "since its a completion call and not chat."
+        )
+        resp = e.responses[0]
+        candidates: list["Candidate"] = resp.candidates
+        if not candidates:
+            msg = str(resp.prompt_feedback)
+        reason_messages = {
+            FinishReason.MAX_TOKENS: (
+                "The maximum number of tokens for the LLM has been reached. Please "
+                "either tweak your prompts or try using another LLM."
+            ),
+            FinishReason.STOP: (
+                "The LLM stopped generating a response due to the natural stop "
+                "point of the model or a provided stop sequence."
+            ),
+            FinishReason.SAFETY: "The LLM response was flagged for safety reasons.",
+            FinishReason.RECITATION: (
+                "The LLM response was flagged for recitation reasons."
+            ),
+            FinishReason.BLOCKLIST: (
+                "The LLM response generation was stopped because it "
+                "contains forbidden terms."
+            ),
+            FinishReason.PROHIBITED_CONTENT: (
+                "The LLM response generation was stopped because it "
+                "potentially contains prohibited content."
+            ),
+            FinishReason.SPII: (
+                "The LLM response generation was stopped because it potentially "
+                "contains Sensitive Personally Identifiable Information."
+            ),
+        }
+
+        err_list = []
+        for candidate in candidates:
+            reason: FinishReason = candidate.finish_reason
+            if candidate.finish_message:
+                err_msg = candidate.finish_message
+            else:
+                err_msg = reason_messages.get(reason, str(candidate))
+            err_list.append(err_msg)
+        msg = "\n\nAnother error: \n".join(err_list)
+        return LLMError(msg)

src/unstract/sdk/llm.py

@@ -11,6 +11,7 @@
 from unstract.sdk.adapter import ToolAdapter
 from unstract.sdk.adapters.constants import Common
 from unstract.sdk.adapters.llm import adapters
+from unstract.sdk.adapters.llm.exceptions import parse_llm_err
 from unstract.sdk.adapters.llm.llm_adapter import LLMAdapter
 from unstract.sdk.constants import LogLevel, ToolEnv
 from unstract.sdk.exceptions import LLMError, RateLimitError, SdkError
@@ -56,9 +57,7 @@ def _initialise(self):
             self._llm_instance = self._get_llm(self._adapter_instance_id)
             self._usage_kwargs["adapter_instance_id"] = self._adapter_instance_id
 
-            if not SdkHelper.is_public_adapter(
-                adapter_id=self._adapter_instance_id
-            ):
+            if not SdkHelper.is_public_adapter(adapter_id=self._adapter_instance_id):
                 platform_api_key = self._tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY)
                 CallbackManager.set_callback(
                     platform_api_key=platform_api_key,
@@ -78,16 +77,8 @@ def complete(
             if match:
                 response.text = match.group(0)
             return {LLM.RESPONSE: response}
-        # TODO: Handle for all LLM providers
-        except OpenAIAPIError as e:
-            msg = "OpenAI error: "
-            if hasattr(e, "body") and "message" in e.body:
-                msg += e.body["message"]
-            else:
-                msg += e.message
-            if isinstance(e, OpenAIRateLimitError):
-                raise RateLimitError(msg)
-            raise LLMError(msg) from e
+        except Exception as e:
+            raise parse_llm_err(e) from e
 
     def _get_llm(self, adapter_instance_id: str) -> LlamaIndexLLM:
         """Returns the LLM object for the tool.

src/unstract/sdk/tool/executor.py

@@ -49,17 +49,22 @@ def execute_run(self, args: argparse.Namespace) -> None:
             self.tool.stream_error_and_exit("--settings are required for RUN command")
         settings: dict[str, Any] = loads(args.settings)
 
-        self._setup_for_run()
+        self.tool.stream_log(
+            f"Running tool with "
+            f"Workflow ID: {self.tool.workflow_id}, "
+            f"Execution ID: {self.tool.execution_id}, "
+            f"SDK Version: {get_sdk_version()}"
+        )
 
+        self._setup_for_run()
         validator = ToolValidator(self.tool)
         settings = validator.validate_pre_execution(settings=settings)
 
         self.tool.stream_log(
-            f"Running tool for "
-            f"Workflow ID: {self.tool.workflow_id}, "
-            f"Execution ID: {self.tool.execution_id}, "
-            f"SDK Version: {get_sdk_version()}, "
+            f"Executing for file: {self.tool.get_exec_metadata['source_name']}, "
+            f"with tool settings: {settings}"
         )
+
         try:
             self.tool.run(
                 settings=settings,

src/unstract/sdk/tool/stream.py

@@ -6,6 +6,7 @@
 from deprecated import deprecated
 
 from unstract.sdk.constants import Command, LogLevel, LogStage, ToolEnv
+from unstract.sdk.utils import ToolUtils
 
 
 class StreamMixin:
@@ -26,7 +27,7 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO, **kwargs) -> None:
 
         """
         self.log_level = log_level
-        self._exec_by_tool = bool(
+        self._exec_by_tool = ToolUtils.str_to_bool(
             os.environ.get(ToolEnv.EXECUTION_BY_TOOL, "False")
         )
         super().__init__(**kwargs)
@@ -78,9 +79,7 @@ def stream_error_and_exit(self, message: str) -> None:
         if self._exec_by_tool:
             exit(1)
         else:
-            raise RuntimeError(
-                "RuntimeError from SDK, check the above log for details"
-            )
+            raise RuntimeError("RuntimeError from SDK, check the above log for details")
 
     def get_env_or_die(self, env_key: str) -> str:
         """Returns the value of an env variable.
@@ -232,9 +231,7 @@ def stream_single_step_message(message: str, **kwargs: Any) -> None:
         print(json.dumps(record))
 
     @staticmethod
-    @deprecated(
-        version="0.4.4", reason="Use `BaseTool.write_to_result()` instead"
-    )
+    @deprecated(version="0.4.4", reason="Use `BaseTool.write_to_result()` instead")
     def stream_result(result: dict[Any, Any], **kwargs: Any) -> None:
         """Streams the result of the tool using the Unstract protocol RESULT to
         stdout.