redesigned context-not-available error to follow the same principles as other IQL errors, inherting from IQLError, thus enabled its handling by self-reflection mechanism

Author: ds-jakub-cierocki

src/dbally/context/exceptions.py

@@ -1,26 +1,5 @@
-from abc import ABC
-
-
-class BaseContextException(Exception, ABC):
-    """
-    A base (abstract) exception for all specification context-related exception.
-    """
-
-
 class ContextNotAvailableError(Exception):
     """
     An exception inheriting from BaseContextException pointining that no sufficient context information
     was provided by the user while calling view.ask().
     """
-
-
-class ContextualisationNotAllowed(Exception):
-    """
-    An exception inheriting from BaseContextException pointining that the filter method signature
-    does not allow to provide an additional context.
-    """
-
-
-# WORKAROUND - traditional inhertiance syntax is not working in context of abstract Exceptions
-BaseContextException.register(ContextNotAvailableError)
-BaseContextException.register(ContextualisationNotAllowed)

src/dbally/iql/_exceptions.py

@@ -1,13 +1,26 @@
 import ast
 from typing import Optional, Union
 
+from typing_extensions import TypeAlias
+
 from dbally.exceptions import DbAllyError
 
+IQLNode: TypeAlias = Union[ast.stmt, ast.expr]
+
 
 class IQLError(DbAllyError):
-    """Base exception for all IQL parsing related exceptions."""
+    """
+    Base exception for all IQL parsing related exceptions.
+
+    Attributes:
+        node: An IQL Node (AST Exprresion) during which processing an error was encountered.
+        source: Raw LLM response containing IQL filter calls.
+    """
+
+    node: IQLNode
+    source: str
 
-    def __init__(self, message: str, node: Union[ast.stmt, ast.expr], source: str) -> None:
+    def __init__(self, message: str, node: IQLNode, source: str) -> None:
         message = message + ": " + source[node.col_offset : node.end_col_offset]
 
         super().__init__(message)
@@ -18,15 +31,15 @@ def __init__(self, message: str, node: Union[ast.stmt, ast.expr], source: str) -
 class IQLArgumentParsingError(IQLError):
     """Raised when an argument cannot be parsed into a valid IQL."""
 
-    def __init__(self, node: Union[ast.stmt, ast.expr], source: str) -> None:
+    def __init__(self, node: IQLNode, source: str) -> None:
         message = "Not a valid IQL argument"
         super().__init__(message, node, source)
 
 
 class IQLUnsupportedSyntaxError(IQLError):
     """Raised when trying to parse an unsupported syntax."""
 
-    def __init__(self, node: Union[ast.stmt, ast.expr], source: str, context: Optional[str] = None) -> None:
+    def __init__(self, node: IQLNode, source: str, context: Optional[str] = None) -> None:
         node_name = node.__class__.__name__
 
         message = f"{node_name} syntax is not supported in IQL"
@@ -47,3 +60,24 @@ def __init__(self, node: ast.Name, source: str) -> None:
 
 class IQLArgumentValidationError(IQLError):
     """Raised when argument is not valid for a given method."""
+
+
+class IQLContextNotAllowedError(IQLError):
+    """
+    Raised when a context call/keyword has been passed as an argument to the filter
+    which does not support contextualization for this specific parameter.
+    """
+
+    def __init__(self, node: IQLNode, source: str, arg_name: Optional[str] = None) -> None:
+        if arg_name is None:
+            message = (
+                "The LLM detected that the context is required to execute the query"
+                "while the filter signature does not allow it at all."
+            )
+        else:
+            message = (
+                "The LLM detected that the context is required to execute the query"
+                f"while the filter signature does allow it for `{arg_name}` argument."
+            )
+
+        super().__init__(message, node, source)

src/dbally/iql/_processor.py

@@ -4,11 +4,11 @@
 from dbally.audit.event_tracker import EventTracker
 from dbally.context._utils import _does_arg_allow_context
 from dbally.context.context import BaseCallerContext, CustomContext
-from dbally.context.exceptions import ContextualisationNotAllowed
 from dbally.iql import syntax
 from dbally.iql._exceptions import (
     IQLArgumentParsingError,
     IQLArgumentValidationError,
+    IQLContextNotAllowedError,
     IQLError,
     IQLFunctionNotExists,
     IQLUnsupportedSyntaxError,
@@ -143,16 +143,10 @@ def _parse_arg(
                 raise IQLArgumentParsingError(arg, self.source)
 
             if parent_func_def.context_class is None:
-                raise ContextualisationNotAllowed(
-                    "The LLM detected that the context is required +\
-                    to execute the query while the filter signature does not allow it at all."
-                )
+                raise IQLContextNotAllowedError(arg, self.source)
 
             if not _does_arg_allow_context(arg_spec):
-                raise ContextualisationNotAllowed(
-                    f"The LLM detected that the context is required +\
-                    to execute the query while the filter signature does allow it for `{arg_spec.name}` argument."
-                )
+                raise IQLContextNotAllowedError(arg, self.source, arg_name=arg_spec.name)
 
             return parent_func_def.context_class.select_context(self.contexts)