nizika_ai/queries_to_answers.py

from __future__ import annotations

import random
from datetime import datetime
from typing import TypedDict

from config import DB
from consts import AnswerType, Character, Platform
from models import Answer, AnsweredFlag, Query, User
from 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:
    if answer.answer_type in [AnswerType.YOUTUBE_REPLY]:
        add_answered_flag (answer, Platform.YOUTUBE)
    if answer.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 ()