core/text2sql/handler.py

import logging
from core.models import MutualFund
from core.text2sql.ml_models import Text2SQLModel

logger = logging.getLogger(__name__)


class QueryDataHandler:
    """
    A class for handling queries and fetching data using a Text2SQL model and Django models.
    """

    def __init__(self):
        self.mutual_fund = MutualFund()
        self.text2sql_model = Text2SQLModel()

    def get_data_from_query(self, prompt):
        """
        Generates a PostgreSQL query using the Text2SQL model based on the input prompt
        and retrieves data using Django models.
        """
        # Use Text2SQL ML model to generate a PostgreSQL query
        sql_query = self.text2sql_model.generate_query(prompt)
        print("SQL Query: ", sql_query)
        logger.info(f"SQL Query: {sql_query}")
        # Use Django models to fetch data
        return sql_query, self.mutual_fund.execute_query(sql_query)