Skip to content

feat: better IQL prompts #98

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: better IQL prompts #98

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6fa0df1
feat: better IQL prompts
ludwiktrammer Apr 16, 2025
234921a
Small prompt fix
ludwiktrammer Apr 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 15 additions & 15 deletions src/dbally/iql_generator/prompt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,17 +194,17 @@ def __init__(
"role": "system",
"content": (
"You have access to an API that lets you query a database:\n"
"\n{methods}\n"
"\n{methods}\n\n"
"Suggest which one(s) to call and how they should be joined with logic operators (AND, OR, NOT).\n"
"Remember! Don't give any comments, just the function calls.\n"
"The output will look like this:\n"
'filter1("arg1") AND (NOT filter2(120) OR filter3(True))\n'
"DO NOT INCLUDE arguments names in your response. Only the values.\n"
"You MUST use only these methods:\n"
"\n{methods}\n"
"It is VERY IMPORTANT not to use methods other than those listed above."
"Remember! Don't return any comments, just the function calls.\n"
"The output should look like Python function calls with positional arguments, joined by logic operators:\n"
'some_filter("foo") AND (NOT filter2(120) OR filter3(True))\n\n'
"DO NOT INCLUDE arguments names in your response. Only the values. Strings must be quoted.\n"
"You MUST use only these functions:\n"
"\n{methods}\n\n"
"It is VERY IMPORTANT not to use functions other than those listed above."
"""If you DON'T KNOW HOW TO ANSWER DON'T SAY anything other than `UNSUPPORTED QUERY`"""
"This is CRUCIAL, otherwise the system will crash. "
"This is CRUCIAL, otherwise the system will crash."
),
},
{
Expand All @@ -220,17 +220,17 @@ def __init__(
{
"role": "system",
"content": (
"You have access to an API that lets you query a database supporting a SINGLE aggregation.\n"
"When prompted for an aggregation, use the following methods: \n"
"You have access to an API that lets you query a database, supporting a SINGLE aggregation.\n"
"When prompted for an aggregation, use one of the following functions: \n"
"{methods}"
"DO NOT INCLUDE arguments names in your response. Only the values.\n"
"You MUST use only these methods:\n"
"The output should look like a single Python function call with positional arguments:\n"
'some_aggregation("foo", True)\n'
"DO NOT INCLUDE arguments names in your response. Only the values. Strings must be quoted.\n\n"
"You MUST only choose from these functions:\n"
"\n{methods}\n"
"It is VERY IMPORTANT not to use methods other than those listed above."
"""If you DON'T KNOW HOW TO ANSWER DON'T SAY anything other than `UNSUPPORTED QUERY`"""
"This is CRUCIAL to put `UNSUPPORTED QUERY` text only, otherwise the system will crash. "
"Structure output to resemble the following pattern:\n"
'aggregation1("arg1", arg2)\n'
),
},
{
Expand Down
72 changes: 46 additions & 26 deletions tests/unit/test_iql_format.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,65 +1,85 @@
from typing import List

from dbally.iql_generator.prompt import FILTERS_GENERATION_TEMPLATE, IQLGenerationPromptFormat
from dbally.prompt.elements import FewShotExample
from dbally.views.exposed_functions import ExposedFunction, MethodParamWithTyping


async def test_iql_prompt_format_default() -> None:
prompt_format = IQLGenerationPromptFormat(
question="",
methods=[],
question="Some question",
methods=[
ExposedFunction(
name="filter_by_name", description="", parameters=[MethodParamWithTyping(name="name", type=List[str])]
),
ExposedFunction(
name="filter_by_city", description="", parameters=[MethodParamWithTyping(name="city", type=str)]
),
ExposedFunction(
name="filter_by_company", description="", parameters=[MethodParamWithTyping(name="company", type=str)]
),
],
examples=[],
)
formatted_prompt = FILTERS_GENERATION_TEMPLATE.format_prompt(prompt_format)

assert formatted_prompt.chat == [
{
"role": "system",
"content": "You have access to an API that lets you query a database:\n"
"\n\n"
"content": "You have access to an API that lets you query a database:\n\n"
+ "filter_by_name(name: List[str])\n"
+ "filter_by_city(city: str)\n"
+ "filter_by_company(company: str)\n\n"
"Suggest which one(s) to call and how they should be joined with logic operators (AND, OR, NOT).\n"
"Remember! Don't give any comments, just the function calls.\n"
"The output will look like this:\n"
'filter1("arg1") AND (NOT filter2(120) OR filter3(True))\n'
"DO NOT INCLUDE arguments names in your response. Only the values.\n"
"You MUST use only these methods:\n"
"\n\n"
"It is VERY IMPORTANT not to use methods other than those listed above."
"Remember! Don't return any comments, just the function calls.\n"
"The output should look like Python function calls with positional arguments, joined by logic operators:\n"
'some_filter("foo") AND (NOT filter2(120) OR filter3(True))\n\n'
"DO NOT INCLUDE arguments names in your response. Only the values. Strings must be quoted.\n"
"You MUST use only these functions:\n\n"
+ "filter_by_name(name: List[str])\n"
+ "filter_by_city(city: str)\n"
+ "filter_by_company(company: str)\n\n"
"It is VERY IMPORTANT not to use functions other than those listed above."
"""If you DON'T KNOW HOW TO ANSWER DON'T SAY anything other than `UNSUPPORTED QUERY`"""
"This is CRUCIAL, otherwise the system will crash. ",
"This is CRUCIAL, otherwise the system will crash.",
"is_example": False,
},
{"role": "user", "content": "", "is_example": False},
{"role": "user", "content": "Some question", "is_example": False},
]


async def test_iql_prompt_format_few_shots_injected() -> None:
examples = [FewShotExample("q1", "a1")]
prompt_format = IQLGenerationPromptFormat(
question="",
methods=[],
question="Some question",
methods=[
ExposedFunction(
name="filter_by_name", description="", parameters=[MethodParamWithTyping(name="name", type=List[str])]
)
],
examples=examples,
)
formatted_prompt = FILTERS_GENERATION_TEMPLATE.format_prompt(prompt_format)

assert formatted_prompt.chat == [
{
"role": "system",
"content": "You have access to an API that lets you query a database:\n"
"\n\n"
"content": "You have access to an API that lets you query a database:\n\n"
+ "filter_by_name(name: List[str])\n\n"
"Suggest which one(s) to call and how they should be joined with logic operators (AND, OR, NOT).\n"
"Remember! Don't give any comments, just the function calls.\n"
"The output will look like this:\n"
'filter1("arg1") AND (NOT filter2(120) OR filter3(True))\n'
"DO NOT INCLUDE arguments names in your response. Only the values.\n"
"You MUST use only these methods:\n"
"\n\n"
"It is VERY IMPORTANT not to use methods other than those listed above."
"Remember! Don't return any comments, just the function calls.\n"
"The output should look like Python function calls with positional arguments, joined by logic operators:\n"
'some_filter("foo") AND (NOT filter2(120) OR filter3(True))\n\n'
"DO NOT INCLUDE arguments names in your response. Only the values. Strings must be quoted.\n"
"You MUST use only these functions:\n\n" + "filter_by_name(name: List[str])\n\n"
"It is VERY IMPORTANT not to use functions other than those listed above."
"""If you DON'T KNOW HOW TO ANSWER DON'T SAY anything other than `UNSUPPORTED QUERY`"""
"This is CRUCIAL, otherwise the system will crash. ",
"This is CRUCIAL, otherwise the system will crash.",
"is_example": False,
},
{"role": "user", "content": examples[0].question, "is_example": True},
{"role": "assistant", "content": examples[0].answer, "is_example": True},
{"role": "user", "content": "", "is_example": False},
{"role": "user", "content": "Some question", "is_example": False},
]


Expand Down
Loading