nizika_ai_service/queries_to_answers.py

"""
DB の queries テーブルにたまってゐるクエリを AI に処理させ answers テーブルに流す．
"""

from __future__ import annotations

import random
from datetime import datetime
from typing import TypedDict

from nizika_ai.config import DB
from nizika_ai.consts import Character, Platform, QueryType
from nizika_ai.models import Answer, AnsweredFlag, Query
from nizika_ai.talk import Talk


def main (
) -> None:
    """
    メーン処理
    """

    queries: list[Query] = Query.where ('answered', False).get ()
    if not queries:
        return

    query: Query = random.choice (queries)

    DB.begin_transaction ()
    try:
        user_name = query.user.name if query.user_id else None

        histories: list[History] = []
        for history in query.answer_histories:
            if history.query is not None:
                histories.append ({ 'role': 'user', 'content': history.query.content })
            histories.append ({ 'role': 'assistant', 'content': history.content })

        for character in [Character.DEERJIKA, Character.GOATOH]:
            if query.target_character & character.value:
                add_answer (query, character, user_name, histories)

        query.answered = True

        query.save ()

        DB.commit ()
    except Exception:
        DB.rollback ()
        raise


def add_answer (
        query:      Query,
        character:  Character,
        user_name:  str | None,
        histories:  list[History],
) -> None:
    """
    AI の返答を DB に積む．

    Parameters
    ----------
    query:      Query
        クエリ
    character:  Character
        返答するキャラクタ
    user_name:  str | None
        クエリの主
    histories:  list[History]
        履歴
    """

    message: str | list[dict[str, str | dict[str, str]]]
    if query.image_url:
        message = [{ 'type': 'text', 'text': query.content },
                   { 'type': 'image_url', 'image_url': query.image_url }]
    else:
        message = query.content

    answer = Answer ()
    answer.query_id = query.id
    answer.character = character.value
    answer.content = Talk.main (message, user_name, histories,
                                goatoh_mode = character == Character.GOATOH)
    answer.sent_at = datetime.now ()
    answer.save ()
    add_answered_flags (answer)


def add_answered_flags (
        answer: Answer,
) -> None:
    """
    返答済フラグを付与する．

    Parameters
    ----------
    answer: Answer
        返答モデル
    """

    answer_type: QueryType
    try:
        answer_type = QueryType (answer.query.query_type)
    except (TypeError, ValueError):
        return

    if answer_type in (QueryType.YOUTUBE_COMMENT,
                       QueryType.YOUTUBE_COMMENT,
                       QueryType.KIRIBAN,
                       QueryType.NICO_REPORT,
                       QueryType.SNACK_TIME,
                       QueryType.HOT_SPRING):
        add_answered_flag (answer, Platform.YOUTUBE)

    if answer_type in (QueryType.BLUESKY_COMMENT,
                       QueryType.BLUESKY_SYSTEM,
                       QueryType.KIRIBAN,
                       QueryType.NICO_REPORT,
                       QueryType.SNACK_TIME,
                       QueryType.HOT_SPRING):
        add_answered_flag (answer, Platform.BLUESKY)


def add_answered_flag (
        answer:     Answer,
        platform:   Platform,
) -> None:
    """
    返答済フラグを付与する．

    Parameters
    ----------
    answer:     Answer
        返答モデル
    platform:   Platform
        プラットフォーム
    """

    answered_flag = AnsweredFlag ()
    answered_flag.answer_id = answer.id
    answered_flag.platform = platform.value
    answered_flag.answered = False
    answered_flag.save ()


class History (TypedDict):
    """
    会話履歴の 1 要素；ユーザや AI の発話を簡易に保持する型
    """

    role:       str
    content:    str


if __name__ == '__main__':
    main ()