diff --git a/src/dbally/collection/collection.py b/src/dbally/collection/collection.py index 5e175cb5..45446cc5 100644 --- a/src/dbally/collection/collection.py +++ b/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() diff --git a/src/dbally/iql/_query.py b/src/dbally/iql/_query.py index 2099d39c..6a610c24 100644 --- a/src/dbally/iql/_query.py +++ b/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) diff --git a/src/dbally/iql_generator/iql_generator.py b/src/dbally/iql_generator/iql_generator.py index 7eeb9154..1946c258 100644 --- a/src/dbally/iql_generator/iql_generator.py +++ b/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)) diff --git a/src/dbally/iql_generator/prompt.py b/src/dbally/iql_generator/prompt.py index 44bb2cd4..0cff10ea 100644 --- a/src/dbally/iql_generator/prompt.py +++ b/src/dbally/iql_generator/prompt.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. " ), diff --git a/src/dbally/views/base.py b/src/dbally/views/base.py index 365e83d6..43b69dbd 100644 --- a/src/dbally/views/base.py +++ b/src/dbally/views/base.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. diff --git a/src/dbally/views/structured.py b/src/dbally/views/structured.py index 5e9c4e1a..eda1a48c 100644 --- a/src/dbally/views/structured.py +++ b/src/dbally/views/structured.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)