feat: request contextualisation - core functionality

src/dbally/collection/collection.py

@@ -3,13 +3,14 @@
 import textwrap
 import time
 from collections import defaultdict
-from typing import Callable, Dict, List, Optional, Type, TypeVar
+from typing import Callable, Dict, Iterable, List, Optional, Type, TypeVar
 
 from dbally.audit.event_handlers.base import EventHandler
 from dbally.audit.event_tracker import EventTracker
 from dbally.audit.events import RequestEnd, RequestStart
 from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
 from dbally.collection.results import ExecutionResult
+from dbally.context.context import CustomContext
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.nl_responder.nl_responder import NLResponder
@@ -156,6 +157,7 @@ async def ask(
         dry_run: bool = False,
         return_natural_response: bool = False,
         llm_options: Optional[LLMOptions] = None,
+        contexts: Optional[Iterable[CustomContext]] = None,
     ) -> ExecutionResult:
         """
         Ask question in a text form and retrieve the answer based on the available views.
@@ -175,6 +177,8 @@ async def ask(
                 the natural response will be included in the answer
             llm_options: options to use for the LLM client. If provided, these options will be merged with the default
                 options provided to the LLM client, prioritizing option values other than NOT_GIVEN
+            contexts: An iterable (typically a list) of context objects, each being an instance of
+                a subclass of BaseCallerContext. May contain contexts irrelevant for the currently processed query.
 
         Returns:
             ExecutionResult object representing the result of the query execution.
@@ -215,6 +219,7 @@ async def ask(
             n_retries=self.n_retries,
             dry_run=dry_run,
             llm_options=llm_options,
+            contexts=contexts,
         )
         end_time_view = time.monotonic()
 

src/dbally/context/__init__.py

@@ -0,0 +1,3 @@
+from .context import BaseCallerContext
+
+__all__ = ["BaseCallerContext"]

src/dbally/context/_utils.py

@@ -0,0 +1,75 @@
+from inspect import isclass
+from typing import Any, Optional, Sequence, Tuple, Type, Union
+
+import typing_extensions as type_ext
+
+from dbally.context.context import BaseCallerContext
+from dbally.views.exposed_functions import MethodParamWithTyping
+
+ContextClass: type_ext.TypeAlias = Optional[Type[BaseCallerContext]]
+
+
+def _extract_params_and_context(
+    filter_method_: type_ext.Callable, hidden_args: Sequence[str]
+) -> Tuple[Sequence[MethodParamWithTyping], ContextClass]:
+    """
+    Processes the MethodsBaseView filter method signauture to extract the args and type hints in the desired format.
+    Context claases are getting excluded the returned MethodParamWithTyping list. Only the first BaseCallerContext
+    class is returned.
+
+    Args:
+        filter_method_: MethodsBaseView filter method (annotated with @decorators.view_filter() decorator)
+        hidden_args: method arguments that should not be extracted
+
+    Returns:
+        The first field contains the list of arguments, each encapsulated as MethodParamWithTyping.
+        The 2nd is the BaseCallerContext subclass provided for this filter, or None if no context specified.
+    """
+
+    params = []
+    context = None
+    # TODO confirm whether to use typing.get_type_hints(method) or method.__annotations__
+    for name_, type_ in type_ext.get_type_hints(filter_method_).items():
+        if name_ in hidden_args:
+            continue
+
+        if isclass(type_) and issubclass(type_, BaseCallerContext):
+            # this is the case when user provides a context but no other type hint for a specifc arg
+            context = type_
+            type_ = Any
+        elif type_ext.get_origin(type_) is Union:
+            union_subtypes = type_ext.get_args(type_)
+            if not union_subtypes:
+                type_ = Any
+
+            for subtype_ in union_subtypes:  # type: ignore
+                # TODO add custom error for the situation when user provides more than two contexts for a single filter
+                # for now we extract only the first context
+                if isclass(subtype_) and issubclass(subtype_, BaseCallerContext):
+                    if context is None:
+                        context = subtype_
+
+        params.append(MethodParamWithTyping(name_, type_))
+
+    return params, context
+
+
+def _does_arg_allow_context(arg: MethodParamWithTyping) -> bool:
+    """
+    Verifies whether a method argument allows contextualization based on the type hints attached to a method signature.
+
+    Args:
+        arg: MethodParamWithTyping container preserving information about the method argument
+
+    Returns:
+        Verification result.
+    """
+
+    if type_ext.get_origin(arg.type) is not Union and not issubclass(arg.type, BaseCallerContext):
+        return False
+
+    for subtype in type_ext.get_args(arg.type):
+        if issubclass(subtype, BaseCallerContext):
+            return True
+
+    return False

src/dbally/context/context.py

@@ -0,0 +1,64 @@
+import ast
+from abc import ABC
+from typing import Iterable
+
+from typing_extensions import Self, TypeAlias
+
+from dbally.context.exceptions import ContextNotAvailableError
+
+CustomContext: TypeAlias = "BaseCallerContext"
+
+
+class BaseCallerContext(ABC):
+    """
+    An interface for contexts that are used to pass additional knowledge about
+    the caller environment to the filters. LLM will always return `Context()`
+    when the context is required and this call will be later substituted by an instance of
+    a class implementing this interface, selected based on the filter method signature (type hints).
+    """
+
+    _alias: str = "Context"
+
+    @classmethod
+    def select_context(cls, contexts: Iterable[CustomContext]) -> Self:
+        """
+        Typically called from a subclass of BaseCallerContext, selects a member object from `contexts` being
+        an instance of the same class. Effectively provides a type dispatch mechanism, substituting the context
+        class by its right instance.
+
+        Args:
+            contexts: A sequence of objects, each being an instance of a different BaseCallerContext subclass.
+
+        Returns:
+            An instance of the same BaseCallerContext subclass this method is caller from.
+
+        Raises:
+            ContextNotAvailableError: If the sequence of context objects passed as argument is empty.
+        """
+
+        if not contexts:
+            raise ContextNotAvailableError(
+                "The LLM detected that the context is required to execute the query +\
+                and the filter signature allows contextualization while the context was not provided."
+            )
+
+        # TODO confirm whether it is possible to design a correct type hints here and skipping `type: ignore`
+        return next(filter(lambda obj: isinstance(obj, cls), contexts))  # type: ignore
+
+    @classmethod
+    def is_context_call(cls, node: ast.expr) -> bool:
+        """
+        Verifies whether an AST node indicates context substitution.
+
+        Args:
+            node: An AST node (expression) to verify:
+
+        Returns:
+            Verification result.
+        """
+
+        return (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id in [cls._alias, cls.__name__]
+        )

src/dbally/context/exceptions.py

@@ -0,0 +1,26 @@
+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/_processor.py

@@ -1,7 +1,10 @@
 import ast
-from typing import TYPE_CHECKING, Any, List, Optional, Union
+from typing import Any, Iterable, List, Mapping, Optional, Union
 
 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,
@@ -11,21 +14,46 @@
     IQLUnsupportedSyntaxError,
 )
 from dbally.iql._type_validators import validate_arg_type
-
-if TYPE_CHECKING:
-    from dbally.views.structured import ExposedFunction
+from dbally.views.exposed_functions import ExposedFunction, MethodParamWithTyping
 
 
 class IQLProcessor:
     """
     Parses IQL string to tree structure.
+
+    Attributes:
+        source: Raw LLM response containing IQL filter calls.
+        allowed_functions: A mapping (typically a dict) of all filters implemented for a certain View.
+        contexts: A sequence (typically a list) of context objects, each being an instance of
+            a subclass of BaseCallerContext. May contain contexts irrelevant for the currently processed query.
     """
 
+    source: str
+    allowed_functions: Mapping[str, "ExposedFunction"]
+    contexts: Iterable[CustomContext]
+    _event_tracker: EventTracker
+
     def __init__(
-        self, source: str, allowed_functions: List["ExposedFunction"], event_tracker: Optional[EventTracker] = None
+        self,
+        source: str,
+        allowed_functions: Iterable[ExposedFunction],
+        contexts: Optional[Iterable[CustomContext]] = None,
+        event_tracker: Optional[EventTracker] = None,
     ) -> None:
+        """
+        IQLProcessor class constructor.
+
+        Args:
+            source: Raw LLM response containing IQL filter calls.
+            allowed_functions: An interable (typically a list) of all filters implemented for a certain View.
+            contexts: An iterable (typically a list) of context objects, each being an instance of
+                a subclass of BaseCallerContext.
+            even_tracker: An EvenTracker instance.
+        """
+
         self.source = source
         self.allowed_functions = {func.name: func for func in allowed_functions}
+        self.contexts = contexts or []
         self._event_tracker = event_tracker or EventTracker()
 
     async def process(self) -> syntax.Node:
@@ -38,6 +66,7 @@ async def process(self) -> syntax.Node:
         Raises:
              IQLError: if parsing fails.
         """
+
         self.source = self._to_lower_except_in_quotes(self.source, ["AND", "OR", "NOT"])
 
         ast_tree = ast.parse(self.source)
@@ -84,13 +113,13 @@ async def _parse_call(self, node: ast.Call) -> syntax.FunctionCall:
         if len(func_def.parameters) != len(node.args):
             raise ValueError(f"The method {func.id} has incorrect number of arguments")
 
-        for arg, arg_def in zip(node.args, func_def.parameters):
-            arg_value = self._parse_arg(arg)
+        for arg, arg_spec in zip(node.args, func_def.parameters):
+            arg_value = self._parse_arg(arg, arg_spec=arg_spec, parent_func_def=func_def)
 
-            if arg_def.similarity_index:
-                arg_value = await arg_def.similarity_index.similar(arg_value, event_tracker=self._event_tracker)
+            if arg_spec.similarity_index:
+                arg_value = await arg_spec.similarity_index.similar(arg_value, event_tracker=self._event_tracker)
 
-            check_result = validate_arg_type(arg_def.type, arg_value)
+            check_result = validate_arg_type(arg_spec.type, arg_value)
 
             if not check_result.valid:
                 raise IQLArgumentValidationError(message=check_result.reason or "", node=arg, source=self.source)
@@ -99,12 +128,37 @@ async def _parse_call(self, node: ast.Call) -> syntax.FunctionCall:
 
         return syntax.FunctionCall(func.id, args)
 
-    def _parse_arg(self, arg: ast.expr) -> Any:
+    def _parse_arg(
+        self,
+        arg: ast.expr,
+        arg_spec: Optional[MethodParamWithTyping] = None,
+        parent_func_def: Optional[ExposedFunction] = None,
+    ) -> Any:
         if isinstance(arg, ast.List):
             return [self._parse_arg(x) for x in arg.elts]
 
+        if BaseCallerContext.is_context_call(arg):
+            if parent_func_def is None or arg_spec is None:
+                # not sure whether this line will be ever reached
+                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."
+                )
+
+            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."
+                )
+
+            return parent_func_def.context_class.select_context(self.contexts)
+
         if not isinstance(arg, ast.Constant):
             raise IQLArgumentParsingError(arg, self.source)
+
         return arg.value
 
     @staticmethod

src/dbally/iql/_query.py

@@ -1,4 +1,8 @@
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Iterable, List, Optional
+
+from typing_extensions import Self
+
+from dbally.context.context import CustomContext
 
 from ..audit.event_tracker import EventTracker
 from . import syntax
@@ -28,16 +32,20 @@ async def parse(
         source: str,
         allowed_functions: List["ExposedFunction"],
         event_tracker: Optional[EventTracker] = None,
-    ) -> "IQLQuery":
+        contexts: Optional[Iterable[CustomContext]] = None,
+    ) -> Self:
         """
         Parse IQL string to IQLQuery object.
 
         Args:
             source: IQL string that needs to be parsed
             allowed_functions: list of IQL functions that are allowed for this query
             event_tracker: EventTracker object to track events
+            contexts: An iterable (typically a list) of context objects, each being
+                an instance of a subclass of BaseCallerContext.
         Returns:
              IQLQuery object
         """
-        root = await IQLProcessor(source, allowed_functions, event_tracker=event_tracker).process()
+
+        root = await IQLProcessor(source, allowed_functions, contexts, event_tracker).process()
         return cls(root=root, source=source)