mlnsio/text2sql

1import logging2from core.models import MutualFund3from core.text2sql.ml_models import Text2SQLModel4 5logger = logging.getLogger(__name__)6 7 8class QueryDataHandler:9    """10    A class for handling queries and fetching data using a Text2SQL model and Django models.11    """12 13    def __init__(self):14        self.mutual_fund = MutualFund()15        self.text2sql_model = Text2SQLModel()16 17    def get_data_from_query(self, prompt):18        """19        Generates a PostgreSQL query using the Text2SQL model based on the input prompt20        and retrieves data using Django models.21        """22        # Use Text2SQL ML model to generate a PostgreSQL query23        sql_query = self.text2sql_model.generate_query(prompt)24        print("SQL Query: ", sql_query)25        logger.info(f"SQL Query: {sql_query}")26        # Use Django models to fetch data27        return sql_query, self.mutual_fund.execute_query(sql_query)28