Add tenancity retries

Author: BenConstable9

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -16,6 +16,7 @@
 import re
 import networkx as nx
 from text_2_sql_core.utils.database import DatabaseEngine
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 logging.basicConfig(level=logging.INFO)
 
@@ -272,6 +273,9 @@ def extract_distinct_values_sql_query(
         """
         return f"""SELECT DISTINCT {column.name} FROM {entity.entity} ORDER BY {column.name} DESC;"""
 
+    @retry(
+        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
+    )
     async def query_entities(
         self, sql_query: str, cast_to: any = None
     ) -> list[EntityItem]:
@@ -565,6 +569,9 @@ async def extract_columns_with_definitions(
 
         return columns
 
+    @retry(
+        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
+    )
     async def send_request_to_llm(self, system_prompt: str, input: str):
         """A method to use GPT to generate a definition for an entity.
 

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/databricks_data_dictionary_creator.py

@@ -6,6 +6,7 @@
 import logging
 import os
 from text_2_sql_core.utils.database import DatabaseEngine
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 
 class DatabricksDataDictionaryCreator(DataDictionaryCreator):
@@ -99,6 +100,9 @@ def extract_entity_relationships_sql_query(self) -> str:
             EntitySchema, Entity, ForeignEntitySchema, ForeignEntity;
         """
 
+    @retry(
+        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
+    )
     async def query_entities(self, sql_query: str, cast_to: any = None) -> list[dict]:
         """
         A method to query a Databricks SQL endpoint for entities.

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/snowflake_data_dictionary_creator.py

@@ -7,6 +7,7 @@
 import os
 
 from text_2_sql_core.utils.database import DatabaseEngine
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 
 class SnowflakeDataDictionaryCreator(DataDictionaryCreator):
@@ -75,6 +76,9 @@ def extract_entity_relationships_sql_query(self) -> str:
             EntitySchema, Entity, ForeignEntitySchema, ForeignEntityConstraint;
         """
 
+    @retry(
+        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
+    )
     async def query_entities(
         self, sql_query: str, cast_to: any = None
     ) -> list[EntityItem]: