ai-meeting-server/app/services/hotword.py

import re
from typing import Any, Optional, AsyncGenerator

from sqlalchemy import select, func, delete

from app.errors.hotword import HotwordOutofLimit, HotwordNotFound
from app.models.db import get_async_session
from app.models.common import Page
from app.models.hotword import (
    ReplaceHotword,
    ASRHotword,
    ASRHotwordModel,
    ReplaceHotwordModel,
)


class ReplaceHotwordService:

    @staticmethod
    async def get_replace_hotword(
        replace_hotword: ReplaceHotwordModel,
    ) -> Optional[ReplaceHotword]:
        async with get_async_session() as session:
            result = await session.execute(
                select(ReplaceHotword).where(ReplaceHotword.id == replace_hotword.id)
            )
            res = result.scalar_one_or_none()
            if res is None:
                raise HotwordNotFound
            return res

    @staticmethod
    async def list_replace_hotwords(page: Page) -> Optional[Any]:
        async with get_async_session() as session:
            result = await session.execute(
                page(select(ReplaceHotword).order_by(ReplaceHotword.id.desc()))
            )
            return result.scalars().all()

    @staticmethod
    async def replace_hotwords_total() -> Optional[Any]:
        async with get_async_session() as session:
            # 计算总记录数
            total_query = select(func.count()).select_from(ReplaceHotword)
            total_result = await session.execute(total_query)
            return total_result.scalar()

    @staticmethod
    async def get_all_replace_hotwords() -> AsyncGenerator[ReplaceHotword, None]:
        async with get_async_session() as session:
            result = await session.execute(select(ReplaceHotword))
            for i in result.scalars():
                yield i

    @staticmethod
    async def create_replace_hotword(
        replace_hotword: ReplaceHotwordModel,
    ) -> int:
        async with get_async_session() as session:
            res = await session.execute(select(func.count()).select_from(ReplaceHotword))
            count = res.scalar_one()
            if count + len(replace_hotword.replace_list) > 200:
                # await session.rollback()
                raise HotwordOutofLimit
            for replace in replace_hotword.replace_list:
                new_replace_hotword = ReplaceHotword(
                    origin_word=replace["origin_word"],
                    hotword=replace["hotword"],
                )
                session.add(new_replace_hotword)

            await session.commit()
            return count + len(replace_hotword.replace_list)

    @staticmethod
    async def update_replace_hotword(
        replace_hotword: ReplaceHotwordModel,
    ) -> Optional[ReplaceHotword]:
        async with get_async_session() as session:
            result = await session.execute(
                select(ReplaceHotword).
                where(ReplaceHotword.id == replace_hotword.id)
            )
            update_replace_hotword = result.scalar_one_or_none()
            if update_replace_hotword is None:
                raise HotwordNotFound
            if replace_hotword.origin_word:
                update_replace_hotword.origin_word = replace_hotword.origin_word
            if replace_hotword.hotword:
                update_replace_hotword.hotword = replace_hotword.hotword
            await session.commit()
            return replace_hotword

    @staticmethod
    async def delete_replace_hotword(
        replace_hotword: ReplaceHotwordModel,
    ) -> Optional[ReplaceHotword]:
        async with get_async_session() as session:
            result = await session.execute(
                select(ReplaceHotword).where(ReplaceHotword.id == replace_hotword.id)
            )
            replace_result = result.scalar_one_or_none()
            if replace_result is None:
                return None
            await session.delete(replace_result)
            await session.commit()
            return None



class ASRHotwordService:

    @staticmethod
    async def get_asr_hotword(asr_hotword: ASRHotwordModel) -> Optional[ASRHotword]:
        async with get_async_session() as session:
            result = await session.execute(
                select(ASRHotword).where(ASRHotword.id == asr_hotword.id)
            )
            res = result.scalar_one_or_none()
            if res is None:
                raise HotwordNotFound
            return res

    @staticmethod
    async def list_asr_hotwords(page: Page) -> Optional[Any]:
        async with get_async_session() as session:
            result = await session.execute(
                page(select(ASRHotword).order_by(ASRHotword.id.desc()))
            )
            return result.scalars().all()

    @staticmethod
    async def asr_hotwords_total() -> Optional[Any]:
        async with get_async_session() as session:
            # 计算总记录数
            total_query = select(func.count()).select_from(ASRHotword)
            total_result = await session.execute(total_query)
            return total_result.scalar()

    @staticmethod
    async def get_all_asr_hotwords() -> AsyncGenerator[ASRHotword, None]:
        async with get_async_session() as session:
            result = await session.execute(select(ASRHotword))
            for i in result.scalars():
                yield i

    @staticmethod
    async def create_asr_hotword(asr_hotword: ASRHotwordModel) -> int:
        async with get_async_session() as session:
            res = await session.execute(select(func.count()).select_from(ASRHotword))
            count = res.scalar_one()
            if count + len(asr_hotword.hotword_list) > 200:
                # await session.rollback()
                raise HotwordOutofLimit
            for hotword in asr_hotword.hotword_list:
                new_hotword = ASRHotword(
                    hotword=hotword
                )
                session.add(new_hotword)

            await session.commit()
            return count + len(asr_hotword.hotword_list)

    @staticmethod
    async def update_asr_hotword(asr_hotword: ASRHotwordModel):
        async with get_async_session() as session:
            result = await session.execute(
                select(ASRHotword).where(ASRHotword.id == asr_hotword.id)
            )
            update_asr_hotword = result.scalar_one_or_none()
            if update_asr_hotword is None:
                raise HotwordNotFound
            if asr_hotword.hotword:
                update_asr_hotword.hotword = asr_hotword.hotword
            if asr_hotword.weight:
                update_asr_hotword.weight = asr_hotword.weight
            await session.commit()

    @staticmethod
    async def delete_asr_hotword(asr_hotword: ASRHotwordModel):
        async with get_async_session() as session:
            res = await session.execute(
                delete(ASRHotword).where(ASRHotword.id == asr_hotword.id)
            )
            await session.commit()
            if res.rowcount == 0:
                raise HotwordNotFound