Skip to content

Commit

Permalink
feat: update amenity_dataset with new hour columns (#151)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Yuan325
Yuan325 authored Jan 8, 2024
1 parent f008c04 commit 160d8df
Show file tree
Hide file tree
Showing 10 changed files with 540 additions and 185 deletions.
325 changes: 160 additions & 165 deletions data/amenity_dataset.csv
View file

Large diffs are not rendered by default.

Empty file added retrieval_service/app/app-test
View file
Empty file.
86 changes: 85 additions & 1 deletion retrieval_service/app/app_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from datetime import datetime
from datetime import datetime, time
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
Expand Down Expand Up @@ -250,6 +250,20 @@ def test_search_airports_with_bad_params(m_datastore, app, params):
terminal="FOO BAR",
category="FEE",
hour="BAZ",
sunday_start_hour=time.fromisoformat("08:00"),
sunday_end_hour=time.fromisoformat("08:00"),
monday_start_hour=time.fromisoformat("08:00"),
monday_end_hour=time.fromisoformat("08:00"),
tuesday_start_hour=time.fromisoformat("08:00"),
tuesday_end_hour=time.fromisoformat("08:00"),
wednesday_start_hour=time.fromisoformat("08:00"),
wednesday_end_hour=time.fromisoformat("08:00"),
thursday_start_hour=time.fromisoformat("08:00"),
thursday_end_hour=time.fromisoformat("08:00"),
friday_start_hour=time.fromisoformat("08:00"),
friday_end_hour=time.fromisoformat("08:00"),
saturday_start_hour=time.fromisoformat("08:00"),
saturday_end_hour=time.fromisoformat("08:00"),
),
{
"id": 1,
Expand All @@ -259,6 +273,20 @@ def test_search_airports_with_bad_params(m_datastore, app, params):
"terminal": "FOO BAR",
"category": "FEE",
"hour": "BAZ",
"sunday_start_hour": "08:00:00",
"sunday_end_hour": "08:00:00",
"monday_start_hour": "08:00:00",
"monday_end_hour": "08:00:00",
"tuesday_start_hour": "08:00:00",
"tuesday_end_hour": "08:00:00",
"wednesday_start_hour": "08:00:00",
"wednesday_end_hour": "08:00:00",
"thursday_start_hour": "08:00:00",
"thursday_end_hour": "08:00:00",
"friday_start_hour": "08:00:00",
"friday_end_hour": "08:00:00",
"saturday_start_hour": "08:00:00",
"saturday_end_hour": "08:00:00",
"content": None,
"embedding": None,
},
Expand Down Expand Up @@ -303,6 +331,20 @@ def test_get_amenity(m_datastore, app, method_name, params, mock_return, expecte
terminal="FOO BAR",
category="FEE",
hour="BAZ",
sunday_start_hour=time.fromisoformat("08:00"),
sunday_end_hour=time.fromisoformat("08:00"),
monday_start_hour=time.fromisoformat("08:00"),
monday_end_hour=time.fromisoformat("08:00"),
tuesday_start_hour=time.fromisoformat("08:00"),
tuesday_end_hour=time.fromisoformat("08:00"),
wednesday_start_hour=time.fromisoformat("08:00"),
wednesday_end_hour=time.fromisoformat("08:00"),
thursday_start_hour=time.fromisoformat("08:00"),
thursday_end_hour=time.fromisoformat("08:00"),
friday_start_hour=time.fromisoformat("08:00"),
friday_end_hour=time.fromisoformat("08:00"),
saturday_start_hour=time.fromisoformat("08:00"),
saturday_end_hour=time.fromisoformat("08:00"),
),
models.Amenity(
id=2,
Expand All @@ -312,6 +354,20 @@ def test_get_amenity(m_datastore, app, method_name, params, mock_return, expecte
terminal="FOO BAR",
category="FEE",
hour="BAZ",
sunday_start_hour=time.fromisoformat("08:00"),
sunday_end_hour=time.fromisoformat("08:00"),
monday_start_hour=time.fromisoformat("08:00"),
monday_end_hour=time.fromisoformat("08:00"),
tuesday_start_hour=time.fromisoformat("08:00"),
tuesday_end_hour=time.fromisoformat("08:00"),
wednesday_start_hour=time.fromisoformat("08:00"),
wednesday_end_hour=time.fromisoformat("08:00"),
thursday_start_hour=time.fromisoformat("08:00"),
thursday_end_hour=time.fromisoformat("08:00"),
friday_start_hour=time.fromisoformat("08:00"),
friday_end_hour=time.fromisoformat("08:00"),
saturday_start_hour=time.fromisoformat("08:00"),
saturday_end_hour=time.fromisoformat("08:00"),
),
],
[
Expand All @@ -323,6 +379,20 @@ def test_get_amenity(m_datastore, app, method_name, params, mock_return, expecte
"terminal": "FOO BAR",
"category": "FEE",
"hour": "BAZ",
"sunday_start_hour": "08:00:00",
"sunday_end_hour": "08:00:00",
"monday_start_hour": "08:00:00",
"monday_end_hour": "08:00:00",
"tuesday_start_hour": "08:00:00",
"tuesday_end_hour": "08:00:00",
"wednesday_start_hour": "08:00:00",
"wednesday_end_hour": "08:00:00",
"thursday_start_hour": "08:00:00",
"thursday_end_hour": "08:00:00",
"friday_start_hour": "08:00:00",
"friday_end_hour": "08:00:00",
"saturday_start_hour": "08:00:00",
"saturday_end_hour": "08:00:00",
"content": None,
"embedding": None,
},
Expand All @@ -334,6 +404,20 @@ def test_get_amenity(m_datastore, app, method_name, params, mock_return, expecte
"terminal": "FOO BAR",
"category": "FEE",
"hour": "BAZ",
"sunday_start_hour": "08:00:00",
"sunday_end_hour": "08:00:00",
"monday_start_hour": "08:00:00",
"monday_end_hour": "08:00:00",
"tuesday_start_hour": "08:00:00",
"tuesday_end_hour": "08:00:00",
"wednesday_start_hour": "08:00:00",
"wednesday_end_hour": "08:00:00",
"thursday_start_hour": "08:00:00",
"thursday_end_hour": "08:00:00",
"friday_start_hour": "08:00:00",
"friday_end_hour": "08:00:00",
"saturday_start_hour": "08:00:00",
"saturday_end_hour": "08:00:00",
"content": None,
"embedding": None,
},
Expand Down
14 changes: 14 additions & 0 deletions retrieval_service/datastore/datastore.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,20 @@ async def export_dataset(
"terminal",
"category",
"hour",
"sunday_start_hour",
"sunday_end_hour",
"monday_start_hour",
"monday_end_hour",
"tuesday_start_hour",
"tuesday_end_hour",
"wednesday_start_hour",
"wednesday_end_hour",
"thursday_start_hour",
"thursday_end_hour",
"friday_start_hour",
"friday_end_hour",
"saturday_start_hour",
"saturday_end_hour",
"content",
"embedding",
]
Expand Down
63 changes: 57 additions & 6 deletions retrieval_service/datastore/providers/cloudsql_postgres.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,20 @@ async def initialize_data(
terminal TEXT,
category TEXT,
hour TEXT,
sunday_start_hour TIME,
sunday_end_hour TIME,
monday_start_hour TIME,
monday_end_hour TIME,
tuesday_start_hour TIME,
tuesday_end_hour TIME,
wednesday_start_hour TIME,
wednesday_end_hour TIME,
thursday_start_hour TIME,
thursday_end_hour TIME,
friday_start_hour TIME,
friday_end_hour TIME,
saturday_start_hour TIME,
saturday_end_hour TIME,
content TEXT NOT NULL,
embedding vector(768) NOT NULL
)
Expand All @@ -140,7 +154,14 @@ async def initialize_data(
# Insert all the data
await conn.execute(
text(
"""INSERT INTO amenities VALUES (:id, :name, :description, :location, :terminal, :category, :hour, :content, :embedding)"""
"""
INSERT INTO amenities VALUES (:id, :name, :description, :location,
:terminal, :category, :hour, :sunday_start_hour, :sunday_end_hour,
:monday_start_hour, :monday_end_hour, :tuesday_start_hour,
:tuesday_end_hour, :wednesday_start_hour, :wednesday_end_hour,
:thursday_start_hour, :thursday_end_hour, :friday_start_hour,
:friday_end_hour, :saturday_start_hour, :saturday_end_hour, :content, :embedding)
"""
),
[
{
Expand All @@ -151,6 +172,20 @@ async def initialize_data(
"terminal": a.terminal,
"category": a.category,
"hour": a.hour,
"sunday_start_hour": a.sunday_start_hour,
"sunday_end_hour": a.sunday_end_hour,
"monday_start_hour": a.monday_start_hour,
"monday_end_hour": a.monday_end_hour,
"tuesday_start_hour": a.tuesday_start_hour,
"tuesday_end_hour": a.tuesday_end_hour,
"wednesday_start_hour": a.wednesday_start_hour,
"wednesday_end_hour": a.wednesday_end_hour,
"thursday_start_hour": a.thursday_start_hour,
"thursday_end_hour": a.thursday_end_hour,
"friday_start_hour": a.friday_start_hour,
"friday_end_hour": a.friday_end_hour,
"saturday_start_hour": a.saturday_start_hour,
"saturday_end_hour": a.saturday_end_hour,
"content": a.content,
"embedding": a.embedding,
}
Expand Down Expand Up @@ -181,7 +216,11 @@ async def initialize_data(
# Insert all the data
await conn.execute(
text(
"""INSERT INTO flights VALUES (:id, :airline, :flight_number, :departure_airport, :arrival_airport, :departure_time, :arrival_time, :departure_gate, :arrival_gate)"""
"""
INSERT INTO flights VALUES (:id, :airline, :flight_number,
:departure_airport, :arrival_airport, :departure_time,
:arrival_time, :departure_gate, :arrival_gate)
"""
),
[
{
Expand Down Expand Up @@ -277,8 +316,12 @@ async def get_amenity(self, id: int) -> Optional[models.Amenity]:
async with self.__pool.connect() as conn:
s = text(
"""
SELECT id, name, description, location, terminal, category, hour
FROM amenities WHERE id=:id
SELECT id, name, description, location, terminal, category, hour,
sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
saturday_start_hour, saturday_end_hour
FROM amenities WHERE id=:id
"""
)
params = {"id": id}
Expand All @@ -296,9 +339,17 @@ async def amenities_search(
async with self.__pool.connect() as conn:
s = text(
"""
SELECT id, name, description, location, terminal, category, hour
SELECT id, name, description, location, terminal, category, hour,
sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
saturday_start_hour, saturday_end_hour
FROM (
SELECT id, name, description, location, terminal, category, hour, 1 - (embedding <=> :query_embedding) AS similarity
SELECT id, name, description, location, terminal, category, hour,
sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
saturday_start_hour, saturday_end_hour, 1 - (embedding <=> :query_embedding) AS similarity
FROM amenities
WHERE 1 - (embedding <=> :query_embedding) > :similarity_threshold
ORDER BY similarity DESC
Expand Down
58 changes: 57 additions & 1 deletion retrieval_service/datastore/providers/cloudsql_postgres_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from datetime import datetime
from datetime import datetime, time
from ipaddress import IPv4Address
from typing import Any, AsyncGenerator, List

Expand Down Expand Up @@ -209,6 +209,20 @@ async def test_get_amenity(ds: cloudsql_postgres.Client):
terminal="Ed Lee International Main Hall",
category="restaurant",
hour="Sunday- Saturday 7:00 am-8:00 pm",
sunday_start_hour=time.fromisoformat("07:00:00"),
sunday_end_hour=time.fromisoformat("20:00:00"),
monday_start_hour=time.fromisoformat("07:00:00"),
monday_end_hour=time.fromisoformat("20:00:00"),
tuesday_start_hour=time.fromisoformat("07:00:00"),
tuesday_end_hour=time.fromisoformat("20:00:00"),
wednesday_start_hour=time.fromisoformat("07:00:00"),
wednesday_end_hour=time.fromisoformat("20:00:00"),
thursday_start_hour=time.fromisoformat("07:00:00"),
thursday_end_hour=time.fromisoformat("20:00:00"),
friday_start_hour=time.fromisoformat("07:00:00"),
friday_end_hour=time.fromisoformat("20:00:00"),
saturday_start_hour=time.fromisoformat("07:00:00"),
saturday_end_hour=time.fromisoformat("20:00:00"),
)
assert res == expected

Expand All @@ -228,6 +242,20 @@ async def test_get_amenity(ds: cloudsql_postgres.Client):
terminal="Ed Lee International Main Hall",
category="restaurant",
hour="Sunday- Saturday 4:00 am-11:00 pm",
sunday_start_hour=time.fromisoformat("04:00:00"),
sunday_end_hour=time.fromisoformat("23:00:00"),
monday_start_hour=time.fromisoformat("04:00:00"),
monday_end_hour=time.fromisoformat("23:00:00"),
tuesday_start_hour=time.fromisoformat("04:00:00"),
tuesday_end_hour=time.fromisoformat("23:00:00"),
wednesday_start_hour=time.fromisoformat("04:00:00"),
wednesday_end_hour=time.fromisoformat("23:00:00"),
thursday_start_hour=time.fromisoformat("04:00:00"),
thursday_end_hour=time.fromisoformat("23:00:00"),
friday_start_hour=time.fromisoformat("04:00:00"),
friday_end_hour=time.fromisoformat("23:00:00"),
saturday_start_hour=time.fromisoformat("04:00:00"),
saturday_end_hour=time.fromisoformat("23:00:00"),
content=None,
embedding=None,
),
Expand All @@ -248,6 +276,20 @@ async def test_get_amenity(ds: cloudsql_postgres.Client):
terminal="International Terminal A",
category="shop",
hour="Sunday - Saturday 8:30 am-11:00 pm",
sunday_start_hour=time.fromisoformat("08:30:00"),
sunday_end_hour=time.fromisoformat("23:00:00"),
monday_start_hour=time.fromisoformat("08:30:00"),
monday_end_hour=time.fromisoformat("23:00:00"),
tuesday_start_hour=time.fromisoformat("08:30:00"),
tuesday_end_hour=time.fromisoformat("23:00:00"),
wednesday_start_hour=time.fromisoformat("08:30:00"),
wednesday_end_hour=time.fromisoformat("23:00:00"),
thursday_start_hour=time.fromisoformat("08:30:00"),
thursday_end_hour=time.fromisoformat("23:00:00"),
friday_start_hour=time.fromisoformat("08:30:00"),
friday_end_hour=time.fromisoformat("23:00:00"),
saturday_start_hour=time.fromisoformat("08:30:00"),
saturday_end_hour=time.fromisoformat("23:00:00"),
content=None,
embedding=None,
),
Expand All @@ -259,6 +301,20 @@ async def test_get_amenity(ds: cloudsql_postgres.Client):
terminal="International Terminal G",
category="shop",
hour="Sunday - Saturday 7:00 am-11:00 pm",
sunday_start_hour=time.fromisoformat("07:00:00"),
sunday_end_hour=time.fromisoformat("23:00:00"),
monday_start_hour=time.fromisoformat("07:00:00"),
monday_end_hour=time.fromisoformat("23:00:00"),
tuesday_start_hour=time.fromisoformat("07:00:00"),
tuesday_end_hour=time.fromisoformat("23:00:00"),
wednesday_start_hour=time.fromisoformat("07:00:00"),
wednesday_end_hour=time.fromisoformat("23:00:00"),
thursday_start_hour=time.fromisoformat("07:00:00"),
thursday_end_hour=time.fromisoformat("23:00:00"),
friday_start_hour=time.fromisoformat("07:00:00"),
friday_end_hour=time.fromisoformat("23:00:00"),
saturday_start_hour=time.fromisoformat("07:00:00"),
saturday_end_hour=time.fromisoformat("23:00:00"),
content=None,
embedding=None,
),
Expand Down
Loading

0 comments on commit 160d8df

Please sign in to comment.