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,
            ),
        )