Dh 4693/column entity checker returns none for big query (#172)

* DH-4693/ColumnEntityChecker returns empty for BigQuery

* DH-4693/lowering the treshold for search result and add like search

dataherald/sql_generator/dataherald_sqlagent.py

@@ -8,6 +8,7 @@
 import numpy as np
 import openai
 import pandas as pd
+import sqlalchemy
 from google.api_core.exceptions import GoogleAPIError
 from langchain.agents.agent import AgentExecutor
 from langchain.agents.agent_toolkits.base import BaseToolkit
@@ -23,7 +24,9 @@
 from langchain.tools.base import BaseTool
 from overrides import override
 from pydantic import BaseModel, Extra, Field
+from sqlalchemy import MetaData
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.sql import func
 
 from dataherald.context_store import ContextStore
 from dataherald.db import DB
@@ -263,12 +266,12 @@ class ColumnEntityChecker(BaseSQLDatabaseTool, BaseTool):
     """
 
     def find_similar_strings(
-        self, input_list: List[tuple], target_string: str, threshold=0.6
+        self, input_list: List[tuple], target_string: str, threshold=0.4
     ):
         similar_strings = []
         for item in input_list:
             similarity = difflib.SequenceMatcher(
-                None, str(item[0]).strip(), target_string
+                None, str(item[0]).strip().lower(), target_string.lower()
             ).ratio()
             if similarity >= threshold:
                 similar_strings.append((str(item[0]).strip(), similarity))
@@ -283,12 +286,28 @@ def _run(
     ) -> str:
         schema, entity = tool_input.split(",")
         table_name, column_name = schema.split("->")
-        query = f"SELECT DISTINCT {column_name} FROM {table_name}"  # noqa: S608
-        results = self.db.run_sql(query)[1]["result"]
+        search_pattern = f"%{entity.strip().lower()}%"
+        meta = MetaData(bind=self.db.engine)
+        table = sqlalchemy.Table(table_name.strip(), meta, autoload=True)
+        search_query = sqlalchemy.select(
+            [func.distinct(table.c[column_name.strip()])]
+        ).where(func.lower(table.c[column_name.strip()]).like(search_pattern))
+        distinct_query = sqlalchemy.select(
+            [func.distinct(table.c[column_name.strip()])]
+        )
+        search_results = self.db.engine.execute(search_query).fetchall()
+        search_results = search_results[:25]
+        results = self.db.engine.execute(distinct_query).fetchall()
         results = self.find_similar_strings(results, entity)
         similar_items = "Similar items:\n"
+        already_added = {}
         for item in results:
             similar_items += f"{item[0]}\n"
+            already_added[item[0]] = True
+        if len(search_results) > 0:
+            for item in search_results:
+                if item[0] not in already_added:
+                    similar_items += f"{item[0]}\n"
         return similar_items
 
     async def _arun(