nizika_ai_service/queries_to_answers.py

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 AnswerType, Character, Platform
from nizika_ai.models import Answer, AnsweredFlag, Query, User
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:
    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.answer_type = query.query_type
    answer.sent_at = datetime.now ()
    answer.save ()
    add_answered_flags (answer)


def add_answered_flags (
        answer: Answer,
) -> None:
    answer_type:    AnswerType
    try:
        answer_type = AnswerType (answer.answer_type)
    except Exception:
        return

    if answer_type in (AnswerType.YOUTUBE_REPLY,):
        add_answered_flag (answer, Platform.YOUTUBE)
    if answer_type in (AnswerType.BLUESKY_REPLY,):
        add_answered_flag (answer, Platform.BLUESKY)


def add_answered_flag (
        answer:     Answer,
        platform:   Platform,
) -> None:
    answered_flag = AnsweredFlag ()
    answered_flag.answer_id = answer.id
    answered_flag.platform = platform.value
    answered_flag.answered = False
    answered_flag.save ()


class History (TypedDict):
    role:       str
    content:    str


if __name__ == '__main__':
    main ()