LLM now pastes Context() instead of BaseCallerContext() to indicate that a context is required

Author: ds-jakub-cierocki

src/dbally/context/context.py

@@ -12,11 +12,13 @@
 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 `BaseCallerContext()`
+    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:
         """
@@ -55,4 +57,8 @@ def is_context_call(cls, node: ast.expr) -> bool:
             Verification result.
         """
 
-        return isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == cls.__name__
+        return (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id in [cls._alias, cls.__name__]
+        )

src/dbally/iql_generator/prompt.py

@@ -74,10 +74,10 @@ def __init__(
                 "You MUST use only these methods:\n"
                 "\n{filters}\n"
                 "It is VERY IMPORTANT not to use methods other than those listed above."
-                "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: BaseCallerContext()."
+                "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: Context()."
                 'The typical input phrase referencing some additional context contains the word "my" or similar phrasing, e.g. "my position name", "my company valuation".'
                 "In that case, the part of the output will look like this:"
-                "filter4(BaseCallerContext())"
+                "filter4(Context())"
                 """If you DON'T KNOW HOW TO ANSWER DON'T SAY \"\", SAY: `UNSUPPORTED QUERY` INSTEAD! """
                 "This is CRUCIAL, otherwise the system will crash. "
             ),

tests/unit/iql/test_iql_parser.py

@@ -26,7 +26,7 @@ async def test_iql_parser():
     custom_context2 = AnotherTestCustomContext(some_field="aaa")
 
     parsed = await IQLQuery.parse(
-        "not (filter_by_name(['John', 'Anne']) and filter_by_city(BaseCallerContext()) and filter_by_company('deepsense.ai'))",
+        "not (filter_by_name(['John', 'Anne']) and filter_by_city(Context()) and filter_by_company('deepsense.ai'))",
         allowed_functions=[
             ExposedFunction(
                 name="filter_by_name", description="", parameters=[MethodParamWithTyping(name="name", type=List[str])]

tests/unit/test_iql_format.py

@@ -23,10 +23,10 @@ async def test_iql_prompt_format_default() -> None:
             "You MUST use only these methods:\n"
             "\n\n"
             "It is VERY IMPORTANT not to use methods other than those listed above."
-            "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: BaseCallerContext()."
+            "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: Context()."
             'The typical input phrase referencing some additional context contains the word "my" or similar phrasing, e.g. "my position name", "my company valuation".'
             "In that case, the part of the output will look like this:"
-            "filter4(BaseCallerContext())"
+            "filter4(Context())"
             """If you DON'T KNOW HOW TO ANSWER DON'T SAY \"\", SAY: `UNSUPPORTED QUERY` INSTEAD! """
             "This is CRUCIAL, otherwise the system will crash. ",
             "is_example": False,
@@ -57,10 +57,10 @@ async def test_iql_prompt_format_few_shots_injected() -> None:
             "You MUST use only these methods:\n"
             "\n\n"
             "It is VERY IMPORTANT not to use methods other than those listed above."
-            "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: BaseCallerContext()."
+            "Finally, if a called function argument value is not directly specified in the query but instead requires knowledge of some additional context, than substitute that argument value by: Context()."
             'The typical input phrase referencing some additional context contains the word "my" or similar phrasing, e.g. "my position name", "my company valuation".'
             "In that case, the part of the output will look like this:"
-            "filter4(BaseCallerContext())"
+            "filter4(Context())"
             """If you DON'T KNOW HOW TO ANSWER DON'T SAY \"\", SAY: `UNSUPPORTED QUERY` INSTEAD! """
             "This is CRUCIAL, otherwise the system will crash. ",
             "is_example": False,