from typing import List

from app.api.routers.request.base import message
from llama_index.core.prompts import PromptTemplate
from llama_index.core.settings import Settings
from pydantic import BaseModel

NEXT_QUESTIONS_SUGGESTION_PROMPT = PromptTemplate(
    "你是一个乐于助人的助手！你的任务是对用户可能会问的下一个问题给出建议。 "
    "\n这是对话历史记录"
    "\n---------------------\n{conversation}\n---------------------"
    "考虑到对话历史记录，仅限于现在知识库已有内容, 请给我 $number_of_questions 个你接下来可能会问题的问题！"
)
N_QUESTION_TO_GENERATE = 3


class NextQuestions(BaseModel):
    """A list of questions that user might ask next"""

    questions: List[str]


class NextQuestionSuggestion:
    @staticmethod
    async def suggest_next_questions(
        message_id: str,
        number_of_questions: int = N_QUESTION_TO_GENERATE,
    ) -> List[str]:
        last_user_message = None
        last_assistant_message = None
        results = message().query(message_id)
        if len(results) > 0:
            last_user_message = results[0]['query']
            last_assistant_message = results[0]['answer']
            conversation: str = f"{last_user_message}\n{last_assistant_message}"
            output: NextQuestions = await Settings.llm.astructured_predict(
                NextQuestions,
                prompt=NEXT_QUESTIONS_SUGGESTION_PROMPT,
                conversation=conversation,
                nun_questions=number_of_questions,
            )
            return output.questions
        return []