Skip to content

Added sql generation and table generation test #55

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
ahujadhruv wants to merge 5 commits into
base: issue-23-Implementing_Test_Cases
Choose a base branch
from
Open

Added sql generation and table generation test #55

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b496e71
Added branch automation and CODEOWNERS (#19)
madhup May 27, 2024
a7121ec
added sql_generation_test.py with initial testing
ahujadhruvgit Jul 8, 2024
19a1b1f
Added optional instead of using | operator
ahujadhruvgit Jul 8, 2024
2c5a5c4
added table_selection_test
ahujadhruvgit Jul 10, 2024
9e0cbbe
added Apache licence
ahujadhruvgit Jul 17, 2024
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
9 changes: 9 additions & 0 deletions .github/CODEOWNERS
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
# This is a comment.
# Each line is a file pattern followed by one or more owners.

# These owners will be the default owners for everything in
# the repo. Unless a later match takes precedence,
# @global-owner1 and @global-owner2 will be requested for
# review when someone opens a pull request.
* @GoogleCloudPlatform/nl2sql-maintainers
/.github/ @GoogleCloudPlatform/nl2sql-maintainers
10 changes: 10 additions & 0 deletions .github/issue-branch.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
autoLinkIssue: true
autoCloseIssue: true
defaultBranch: 'dev'
openDraftPR: true
copyIssueDescriptionToPR: true
copyIssueLabelsToPR: true
copyIssueAssigneeToPR: true
copyIssueProjectsToPR: true
copyIssueMilestoneToPR: true
conventionalPrTitles: true
15 changes: 15 additions & 0 deletions .github/workflows/issue-branch.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
on:
issues:
types: [ assigned ]
# The pull_request events below are only needed for pull-request related features.
pull_request:
types: [ opened, closed ]

jobs:
create_issue_branch_job:
runs-on: ubuntu-latest
steps:
- name: Create Issue Branch
uses: robvanderleek/create-issue-branch@main
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
6 changes: 3 additions & 3 deletions nl2sql/tasks/sql_generation/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from loguru import logger
from pydantic import BaseModel, SkipValidation
from typing_extensions import Literal

from typing import Optional
from nl2sql.assets.prompts import FewShot as FewShotPrompts
from nl2sql.assets.prompts import ZeroShot as ZeroShotPrompts
from nl2sql.datasets.base import Database
Expand Down Expand Up @@ -115,9 +115,9 @@ def LANGCHAIN_ZERO_SHOT_PROMPT(self) -> _CoreSqlGeneratorPrompt:
def custom_prompt(
cls,
prompt_template: BasePromptTemplate,
parser: StructuredOutputParser | None = None,
parser: Optional[StructuredOutputParser] = None,
post_processor: Callable = lambda x: x,
prompt_template_id: str | None = None,
prompt_template_id: Optional[str] = None,
) -> _CoreSqlGeneratorPrompt:
"""
Use a custom PromptTemplate for SQL Generation.
Expand Down
52 changes: 52 additions & 0 deletions tests/tasks/sql_generation_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from unittest.mock import MagicMock
from nl2sql.tasks.sql_generation.core import CoreSqlGenerator, CoreSqlGenratorResult

class TestCoreSqlGenerator(unittest.TestCase):
def test_core_sql_generator_with_valid_response(self):
mock_llm = MagicMock()
mock_llm.generate.return_value = MagicMock(
generations=[
[
MagicMock(text="SELECT AVG(price) FROM products WHERE category = 'Electronics';")
]
]
)

mock_db = MagicMock()
mock_db.db.dialect = "sqlite"
mock_db.db.table_info = {
"products": {"product_id": "INT PRIMARY KEY", "name": "TEXT", "price": "REAL", "category": "TEXT"}
}
mock_db.db._usable_tables = ["products"]
mock_db.name = "test_db"
mock_db.descriptor = "A test database"

# Initialize with the mock LLM
generator = CoreSqlGenerator(llm=mock_llm)

# Run the generator
result = generator(mock_db, "What is the average price of products in the 'Electronics' category?")

# Assertions
self.assertEqual(result.generated_query, "SELECT AVG(price) FROM products WHERE category = 'Electronics';")
self.assertEqual(result.db_name, "test_db")
self.assertEqual(result.question, "What is the average price of products in the 'Electronics' category?")
self.assertEqual(len(result.intermediate_steps), 1)

# Verify LLM call
mock_llm.generate.assert_called_once()
112 changes: 112 additions & 0 deletions tests/tasks/table_selection_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from unittest.mock import MagicMock
from loguru import logger
from nl2sql.datasets.base import Database
from nl2sql.tasks.table_selection.core import (
CoreTableSelector,
_TableSelectorPrompts,
)

class TestCoreTableSelector(unittest.TestCase):
def test_call_with_langchain_decider_prompt(self):
mock_llm = MagicMock()
mock_llm.generate.return_value = MagicMock(
generations=[
[MagicMock(text="TableA, TableB")]
]
)
selector = CoreTableSelector(
llm=mock_llm, prompt=_TableSelectorPrompts().LANGCHAIN_DECIDER_PROMPT
)
mock_db = MagicMock()
mock_db.name = "test_db"
mock_db.db._usable_tables = {"TableA", "TableC"}
mock_db.descriptor = {
"TableA": "Description A",
"TableB": "Description B",
"TableC": "Description C",
}
result = selector(mock_db, "test question")
self.assertEqual(result.selected_tables, {"TableA"})
self.assertEqual(result.db_name, "test_db")
self.assertEqual(result.question, "test question")
self.assertEqual(result.available_tables, {"TableA", "TableC"})

def test_call_with_curated_few_shot_cot_prompt(self):
mock_llm = MagicMock()
mock_llm.generate.side_effect = [
MagicMock(
generations=[
[MagicMock(text="Yes. TableA is relevant")]
]
),
MagicMock(
generations=[
[MagicMock(text="No. TableB is not relevant")]
]
),
]
selector = CoreTableSelector(
llm=mock_llm, prompt=_TableSelectorPrompts().CURATED_FEW_SHOT_COT_PROMPT
)
mock_db = MagicMock()
mock_db.name = "test_db"
mock_db.db._usable_tables = {"TableA", "TableB"}
mock_db.descriptor = {
"TableA": {
"col_descriptor": {
"column1": "data_type",
"column2": "data_type"
}
},
"TableB": {
"col_descriptor": {
"column1": "data_type",
"column2": "data_type"
}
}
}
result = selector(mock_db, "test question")
self.assertEqual(result.selected_tables, {"TableA"})
self.assertEqual(result.db_name, "test_db")
self.assertEqual(result.question, "test question")
self.assertEqual(result.available_tables, {"TableA", "TableB"})

def test_call_with_empty_response(self):
mock_llm = MagicMock()
mock_llm.generate.return_value = MagicMock(
generations=[
[MagicMock(text=" ")] # Set text to an empty string
]
)
selector = CoreTableSelector(
llm=mock_llm, prompt=_TableSelectorPrompts().LANGCHAIN_DECIDER_PROMPT
)
mock_db = MagicMock()
mock_db.name = "test_db"
mock_db.db._usable_tables = {"TableA", "TableC"}
mock_db.descriptor = {
"TableA": "Description A",
"TableB": "Description B",
"TableC": "Description C",
}
# with self.assertLogs("nl2sql.tasks.table_selection.core", level="CRITICAL"):
result = selector(mock_db, "test question")
self.assertEqual(result.selected_tables, set())
self.assertEqual(result.db_name, "test_db")
self.assertEqual(result.question, "test question")
self.assertEqual(result.available_tables, {"TableA", "TableC"})