diff --git a/parser_api/controller.py b/parser_api/controller.py index 3c375b0..cab2336 100644 --- a/parser_api/controller.py +++ b/parser_api/controller.py @@ -8,6 +8,8 @@ from sqlalchemy.orm import Session from sqlalchemy.sql import operators from sqlalchemy.sql.expression import BinaryExpression +from rosseti_parser import get_building_id + from . import models, schemas @@ -81,6 +83,11 @@ def search_all(db: Session, prompt: str) -> List[schemas.Record]: 'res' ), contains_lower('comment', prompt)) + building_id, *_ = get_building_id(prompt) + + if building_id is not None: + query |= models.Record.building_id == building_id + res = db.query(models.Record).filter(query).all() return res diff --git a/rosseti_parser/README.md b/rosseti_parser/README.md index 033af96..5672c3d 100644 --- a/rosseti_parser/README.md +++ b/rosseti_parser/README.md @@ -29,6 +29,10 @@ class RossetiParser: ```python def split_addresses(df: pd.DataFrame) -> pd.DataFrame ``` +- `get_building_id`: +```python +def get_building_id(street: str) -> Tuple[Optional[int], Optional[float], Optional[float]] +``` - `fetch_builing_ids`: ```python def fetch_builing_ids(df: pd.DataFrame) -> pd.DataFrame diff --git a/rosseti_parser/__init__.py b/rosseti_parser/__init__.py index fe5e7e9..402d19c 100644 --- a/rosseti_parser/__init__.py +++ b/rosseti_parser/__init__.py @@ -1,5 +1,5 @@ from .rosseti import RossetiParser from .address import split_addresses -from .building_id import fetch_builing_ids, async_fetch_building_ids, concurrent_fetch_builing_ids +from .building_id import get_building_id, fetch_builing_ids, async_fetch_building_ids, concurrent_fetch_builing_ids from .preprocess import preprocess_df, COL_NS, ICOL_NS, preprocess_read_df, group_by_index from .util import pipeline