from datetime import date, datetime
from decimal import Decimal

from langchain.chains import LLMChain
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from sqlalchemy import text

from dataherald.model.chat_model import ChatModel
from dataherald.repositories.database_connections import DatabaseConnectionRepository
from dataherald.repositories.prompts import PromptRepository
from dataherald.sql_database.base import SQLDatabase, SQLInjectionError
from dataherald.types import LLMConfig, NLGeneration, SQLGeneration

HUMAN_TEMPLATE = """Given a Question, a Sql query and the sql query result try to answer the question
If the sql query result doesn't answer the question just say 'I don't know'
Answer the question given the sql query and the sql query result.
Question: {prompt}
SQL query: {sql_query}
SQL query result: {sql_query_result}
"""


class GeneratesNlAnswer:
    def __init__(self, system, storage, llm_config: LLMConfig):
        self.system = system
        self.storage = storage
        self.llm_config = llm_config
        self.model = ChatModel(self.system)

    def execute(
        self,
        sql_generation: SQLGeneration,
        top_k: int = 100,
    ) -> NLGeneration:
        prompt_repository = PromptRepository(self.storage)
        prompt = prompt_repository.find_by_id(sql_generation.prompt_id)

        db_connection_repository = DatabaseConnectionRepository(self.storage)
        database_connection = db_connection_repository.find_by_id(
            prompt.db_connection_id
        )
        self.llm = self.model.get_model(
            database_connection=database_connection,
            temperature=0,
            model_name=self.llm_config.llm_name,
            api_base=self.llm_config.api_base,
        )
        database = SQLDatabase.get_sql_engine(database_connection, True)

        if sql_generation.status == "INVALID":
            return NLGeneration(
                sql_generation_id=sql_generation.id,
                text="I don't know, the SQL query is invalid.",
                created_at=datetime.now(),
            )

        try:
            query = database.parser_to_filter_commands(sql_generation.sql)
            with database._engine.connect() as connection:
                execution = connection.execute(text(query))
                result = execution.fetchmany(top_k)
            rows = []
            for row in result:
                modified_row = {}
                for key, value in zip(row.keys(), row, strict=True):
                    if type(value) in [
                        date,
                        datetime,
                    ]:  # Check if the value is an instance of datetime.date
                        modified_row[key] = str(value)
                    elif (
                        type(value) is Decimal
                    ):  # Check if the value is an instance of decimal.Decimal
                        modified_row[key] = float(value)
                    else:
                        modified_row[key] = value
                rows.append(modified_row)

        except SQLInjectionError as e:
            raise SQLInjectionError(
                "Sensitive SQL keyword detected in the query."
            ) from e

        human_message_prompt = HumanMessagePromptTemplate.from_template(HUMAN_TEMPLATE)
        chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])
        chain = LLMChain(llm=self.llm, prompt=chat_prompt)
        nl_resp = chain.invoke(
            {
                "prompt": prompt.text,
                "sql_query": sql_generation.sql,
                "sql_query_result": "\n".join([str(row) for row in rows]),
            }
        )
        return NLGeneration(
            sql_generation_id=sql_generation.id,
            llm_config=self.llm_config,
            text=nl_resp["text"],
            created_at=datetime.now(),
        )

langchain.prompts.chat.HumanMessagePromptTemplate.from_template langchain.chains.LLMChain langchain.prompts.chat.ChatPromptTemplate.from_messages
[((1101, 1123), 'dataherald.model.chat_model.ChatModel', 'ChatModel', (['self.system'], {}), '(self.system)\n', (1110, 1123), False, 'from dataherald.model.chat_model import ChatModel\n'), ((1272, 1302), 'dataherald.repositories.prompts.PromptRepository', 'PromptRepository', (['self.storage'], {}), '(self.storage)\n', (1288, 1302), False, 'from dataherald.repositories.prompts import PromptRepository\n'), ((1411, 1453), 'dataherald.repositories.database_connections.DatabaseConnectionRepository', 'DatabaseConnectionRepository', (['self.storage'], {}), '(self.storage)\n', (1439, 1453), False, 'from dataherald.repositories.database_connections import DatabaseConnectionRepository\n'), ((1813, 1866), 'dataherald.sql_database.base.SQLDatabase.get_sql_engine', 'SQLDatabase.get_sql_engine', (['database_connection', '(True)'], {}), '(database_connection, True)\n', (1839, 1866), False, 'from dataherald.sql_database.base import SQLDatabase, SQLInjectionError\n'), ((3295, 3351), 'langchain.prompts.chat.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['HUMAN_TEMPLATE'], {}), '(HUMAN_TEMPLATE)\n', (3335, 3351), False, 'from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate\n'), ((3374, 3430), 'langchain.prompts.chat.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (['[human_message_prompt]'], {}), '([human_message_prompt])\n', (3406, 3430), False, 'from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate\n'), ((3447, 3489), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'self.llm', 'prompt': 'chat_prompt'}), '(llm=self.llm, prompt=chat_prompt)\n', (3455, 3489), False, 'from langchain.chains import LLMChain\n'), ((3160, 3225), 'dataherald.sql_database.base.SQLInjectionError', 'SQLInjectionError', (['"""Sensitive SQL keyword detected in the query."""'], {}), "('Sensitive SQL keyword detected in the query.')\n", (3177, 3225), False, 'from dataherald.sql_database.base import SQLDatabase, SQLInjectionError\n'), ((3898, 3912), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (3910, 3912), False, 'from datetime import date, datetime\n'), ((2092, 2106), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (2104, 2106), False, 'from datetime import date, datetime\n'), ((2317, 2328), 'sqlalchemy.text', 'text', (['query'], {}), '(query)\n', (2321, 2328), False, 'from sqlalchemy import text\n')]