post-merge fixes + minor refactor

Author: ds-jakub-cierocki

src/dbally/collection/collection.py

@@ -157,7 +157,7 @@ async def ask(
         dry_run: bool = False,
         return_natural_response: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[CustomContextsList] = None
+        contexts: Optional[CustomContextsList] = None
     ) -> ExecutionResult:
         """
         Ask question in a text form and retrieve the answer based on the available views.
@@ -217,7 +217,7 @@ async def ask(
             n_retries=self.n_retries,
             dry_run=dry_run,
             llm_options=llm_options,
-            context=context
+            contexts=contexts
         )
         end_time_view = time.monotonic()
 

src/dbally/iql/_query.py

@@ -30,7 +30,7 @@ async def parse(
         source: str,
         allowed_functions: List["ExposedFunction"],
         event_tracker: Optional[EventTracker] = None,
-        context: Optional[CustomContextsList] = None
+        contexts: Optional[CustomContextsList] = None
     ) -> Self:
         """
         Parse IQL string to IQLQuery object.
@@ -43,5 +43,5 @@ async def parse(
              IQLQuery object
         """
 
-        root = await IQLProcessor(source, allowed_functions, context, event_tracker).process()
+        root = await IQLProcessor(source, allowed_functions, contexts, event_tracker).process()
         return cls(root=root, source=source)

src/dbally/iql_generator/iql_generator.py

@@ -8,6 +8,7 @@
 from dbally.prompt.elements import FewShotExample
 from dbally.prompt.template import PromptTemplate
 from dbally.views.exposed_functions import ExposedFunction
+from dbally.context.context import CustomContextsList
 
 ERROR_MESSAGE = "Unfortunately, generated IQL is not valid. Please try again, \
     generation of correct IQL is very important. Below you have errors generated by the system:\n{error}"
@@ -42,7 +43,8 @@ async def generate_iql(
         examples: Optional[List[FewShotExample]] = None,
         llm_options: Optional[LLMOptions] = None,
         n_retries: int = 3,
-    ) -> IQLQuery:
+        contexts: Optional[CustomContextsList] = None
+    ) -> Optional[IQLQuery]:
         """
         Generates IQL in text form using LLM.
 
@@ -60,7 +62,7 @@ async def generate_iql(
         prompt_format = IQLGenerationPromptFormat(
             question=question,
             filters=filters,
-            examples=examples,
+            examples=examples or [],
         )
         formatted_prompt = self._prompt_template.format_prompt(prompt_format)
 
@@ -78,7 +80,9 @@ async def generate_iql(
                     source=iql,
                     allowed_functions=filters,
                     event_tracker=event_tracker,
+                    contexts=contexts
                 )
             except IQLError as exc:
+                # TODO handle the possibility of variable `response` being not initialized while runnning the following line
                 formatted_prompt = formatted_prompt.add_assistant_message(response)
                 formatted_prompt = formatted_prompt.add_user_message(ERROR_MESSAGE.format(error=exc))

src/dbally/iql_generator/iql_prompt_template.py

@@ -74,6 +74,11 @@ 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."
+                "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()."
+                '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())"
+                "It is VERY IMPORTANT not to use methods other than those listed above."
                 """If you DON'T KNOW HOW TO ANSWER DON'T SAY \"\", SAY: `UNSUPPORTED QUERY` INSTEAD! """
                 "This is CRUCIAL, otherwise the system will crash. "
             ),

src/dbally/iql_generator/prompt.py

@@ -27,7 +27,7 @@ async def ask(
         n_retries: int = 3,
         dry_run: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[CustomContextsList] = None
+        contexts: Optional[CustomContextsList] = None
     ) -> ViewExecutionResult:
         """
         Executes the query and returns the result.

src/dbally/views/base.py

@@ -42,7 +42,7 @@ async def ask(
         n_retries: int = 3,
         dry_run: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[CustomContextsList] = None
+        contexts: Optional[CustomContextsList] = None
     ) -> ViewExecutionResult:
         """
         Executes the query and returns the result. It generates the IQL query from the natural language query\
@@ -71,6 +71,7 @@ async def ask(
             event_tracker=event_tracker,
             llm_options=llm_options,
             n_retries=n_retries,
+            contexts=contexts
         )
 
         await self.apply_filters(iql)