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:
    DB.begin_transaction ()
    queries: list[Query] = Query.where ('answered', False).get ()
    if not queries:
        return
    query: Query = random.choice (queries)
    user: User = query.user
    user_name: str | None = None
    if query.user_id is not None:
        user_name = user.name
    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 ()


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 is None:
        message = query.content
    else:
        message = [{ 'type': 'text', 'text': query.content },
                   { 'type': 'image_url', 'image_url': query.image_url }]
    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 ()