Spaces:
Running
Running
| # from src import * | |
| import logging | |
| import os | |
| import lancedb | |
| from lancedb.rerankers import ColbertReranker | |
| import sys | |
| logger = logging.getLogger(__name__) | |
| from typing import Optional | |
| from src.vectordb.helpers import set_uri | |
| # db = lancedb.connect("/tmp/db") | |
| def search(query: str, table_name: str, filter_condition: Optional[str] = None, | |
| category: str = "docs", limit: int = 10, reranking: int = 0, | |
| run_local: Optional[bool] = False) -> list | None: | |
| """ | |
| Generalized function to search a database table, with optional filters and reranking. | |
| Args: | |
| - query: str, search query. | |
| - table_name: str, name of the table to search. | |
| - filter_condition: Optional[str], optional SQL-like condition for filtering results. | |
| - category: str, type of category (default is 'docs'). | |
| - limit: int, number of results (default is 10). | |
| - reranking: int (0 or 1), if activated, ColbertReranker is used. | |
| - run_local: Optional[bool], whether to run in a local environment. | |
| Returns: | |
| A list of the most relevant documents or listings based on the category. | |
| """ | |
| uri = set_uri(run_local) | |
| try: | |
| db = lancedb.connect(uri) | |
| except Exception as e: | |
| logger.error(f"Error while connecting to DB: {e}") | |
| return None | |
| logger.info(f"Connected to {table_name} DB.") | |
| table = db.open_table(table_name) | |
| search_query = table.search(query).metric('cosine') | |
| if filter_condition: | |
| search_query = search_query.where(filter_condition) | |
| if reranking: | |
| try: | |
| column = 'description' if category == 'listings' else 'text' | |
| reranker = ColbertReranker(column=column) | |
| results = search_query.rerank(reranker=reranker).limit(limit).to_list() | |
| except Exception as e: | |
| exc_type, exc_obj, exc_tb = sys.exc_info() | |
| fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] | |
| logger.error(f"Error while reranking results: {e}, {(exc_type, fname, exc_tb.tb_lineno)}") | |
| return None | |
| else: | |
| try: | |
| results = search_query.limit(limit).to_list() | |
| except Exception as e: | |
| exc_type, exc_obj, exc_tb = sys.exc_info() | |
| fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] | |
| logger.error(f"Error while searching: {e}, {(exc_type, fname, exc_tb.tb_lineno)}") | |
| return None | |
| logger.info("Found the most relevant documents.") | |
| if category == "docs": | |
| return [{"city": r['city'], "country": r['country'], "section": r['section'], "text": r['text']} for r in | |
| results] | |
| else: | |
| return [{"city": r['city'], "country": r['country'], "type": r['type'], "title": r['title'], | |
| "description": r['description']} for r in results] | |
| def search_wikivoyage_docs(query: str, limit: int = 10, reranking: int = 0, | |
| run_local: Optional[bool] = False) -> list | None: | |
| """ | |
| Function to search documents in the Wikivoyage database. | |
| """ | |
| return search(query=query, table_name="wikivoyage_documents", category="docs", | |
| limit=limit, reranking=reranking, run_local=run_local) | |
| def search_wikivoyage_listings(query: str, cities: list, limit: int = 10, reranking: int = 0, | |
| run_local: Optional[bool] = False) -> list | None: | |
| """ | |
| Function to search listings in the Wikivoyage database, post-filtered by cities. | |
| """ | |
| cities_filter = f"city IN {tuple(cities)}" | |
| return search(query=query, table_name="wikivoyage_listings", filter_condition=cities_filter, | |
| category="listings", limit=limit, reranking=reranking, run_local=run_local) | |