rag_server/applications/utils/mysql/mapper.py

from .base import BaseMySQLClient


class Dataset(BaseMySQLClient):
    async def update_dataset_status(self, dataset_id, ori_status, new_status):
        query = """
            UPDATE dataset SET status = %s WHERE id = %s AND status = %s;
        """
        return await self.pool.async_save(
            query=query, params=(new_status, dataset_id, ori_status)
        )

    async def select_dataset(self, status=1):
        query = """
            SELECT * FROM dataset WHERE status = %s;
        """
        return await self.pool.async_fetch(query=query, params=(status,))

    async def add_dataset(self, name):
        query = """
            INSERT INTO dataset (name) VALUES (%s);
        """
        return await self.pool.async_save(query=query, params=(name,))

    async def select_dataset_by_id(self, id_, status: int = 1):
        query = """
            SELECT * FROM dataset WHERE id = %s AND status = %s;
        """
        return await self.pool.async_fetch(query=query, params=(id_, status))

    async def select_dataset_by_name(self, name, status: int = 1):
        query = """
            SELECT * FROM dataset WHERE name = %s AND status = %s;
        """
        return await self.pool.async_fetch(query=query, params=(name, status))


class ChatResult(BaseMySQLClient):
    async def insert_chat_result(
        self, query_text, dataset_ids, search_res, chat_res, score, has_answer, ai_answer, ai_source, ai_status
    ):
        query = """
                    INSERT INTO chat_res
                        (query, dataset_ids, search_res, chat_res, score, has_answer, ai_answer, ai_source, ai_status) 
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s);
                """
        return await self.pool.async_save(
            query=query,
            params=(query_text, dataset_ids, search_res, chat_res, score, has_answer, ai_answer, ai_source, ai_status),
        )