# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/musique.qa.ipynb.

# %% auto 0
__all__ = ['log', 'DEFAULT_MODEL', 'DEFAULT_COMPLETION_KWARGS', 'DEFAULT_USER_PROMPT_TEMPLATE', 'SYSTEM_PROMPT_STANDARD',
           'answer_question_standard', 'SYSTEM_PROMPT_COT', 'answer_question_cot_zs', 'FEW_SHOT_EXAMPLES_COT',
           'answer_question_cot_fs', 'SYSTEM_PROMPT_CTE', 'answer_question_cte_zs', 'FEW_SHOT_EXAMPLES_CTE',
           'answer_question_cte_fs', 'make_qa_func', 'load_qa_func']

# %% ../../nbs/musique.qa.ipynb 4
from typing import Callable

import openai

from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/musique.qa.ipynb 5
DEFAULT_MODEL = "gpt-3.5-turbo"
DEFAULT_COMPLETION_KWARGS = {"temperature": 0.1}

# %% ../../nbs/musique.qa.ipynb 7
DEFAULT_USER_PROMPT_TEMPLATE = """The context information is provided below.
---------------------
{context}
---------------------
Given the context information and not prior knowledge, answer the question.
{question}
"""


def make_qa_func(
    system_prompt: str,
    user_prompt_template: str = DEFAULT_USER_PROMPT_TEMPLATE,
    few_shot_examples: list[dict] | None = None,
):
    few_shot_examples = few_shot_examples or []

    def answer_question(
        context: str,
        question: str,
        model_name: str = DEFAULT_MODEL,
        completion_kwargs: dict | None = None,
        client=None,
    ) -> dict:
        if client is None:
            client = openai.Client()

        if completion_kwargs is None:
            completion_kwargs = DEFAULT_COMPLETION_KWARGS

        # Prepare the messages
        few_shot_messages = []
        for example in few_shot_examples:
            few_shot_messages.extend(
                [
                    {
                        "role": "user",
                        "content": user_prompt_template.format(
                            context=example["context"],
                            question=example["question"],
                        ),
                    },
                    {
                        "role": "assistant",
                        "content": example["generation"],
                    },
                ]
            )
        messages = [
            {"role": "system", "content": system_prompt},
            *few_shot_messages,
            {
                "role": "user",
                "content": user_prompt_template.format(context=context, question=question),
            },
        ]
        chat_completion = client.chat.completions.create(
            model=model_name,
            messages=messages,
            **completion_kwargs,
        )
        generation = chat_completion.choices[0].message.content
        parts = generation.split("Answer:")
        if len(parts) < 2:
            return dict(answer="", generation=generation)
        answer = parts[1].strip()
        return dict(answer=answer, generation=generation)

    return answer_question

# %% ../../nbs/musique.qa.ipynb 9
SYSTEM_PROMPT_STANDARD = """
You are an excellent question-answering system known for providing accurate and reliable answers. Your responses should be solely based on the context information given, without drawing on prior knowledge. 

# Output format
Answer: [answer in least number of words possible]
""".strip()


answer_question_standard = make_qa_func(
    system_prompt=SYSTEM_PROMPT_STANDARD,
)

# %% ../../nbs/musique.qa.ipynb 12
SYSTEM_PROMPT_COT = """You are an excellent question-answering system known for providing accurate and reliable answers. Your responses should be solely based on the context information given, without drawing on prior knowledge. Always provide clear and logical step-by-step reasoning in your response.

# Output format
Reasoning: [Step-by-step reasoning for the answer.]
Answer: [answer in least number of words possible]
"""

answer_question_cot_zs = make_qa_func(
    system_prompt=SYSTEM_PROMPT_COT,
)

# %% ../../nbs/musique.qa.ipynb 14
FEW_SHOT_EXAMPLES_COT = [
    {
        "id": "2hop__784447_126070",
        "context": 'Glenhis Hern\u00e1ndez (born 7 October 1990 in Havana) is a taekwondo practitioner from Cuba. She was the 2013 World\nChampion in middleweight.\n\nThe current mayor of Havana ("President of the People\'s Power Provincial Assembly") is Marta Hern\u00e1ndez Romero, she\nwas elected on March 5, 2011.',
        "question": "Who is the current mayor of the city Glenhis Hern\u00e1ndez was born?",
        "generation": "Reasoning:\n1. Glenhis Hernández was born in Havana, as mentioned in the context.\n2. The current mayor of Havana mentioned in the context is Marta Hernández Romero.\n3. Therefore, the current mayor of the city where Glenhis Hernández was born is Marta Hernández Romero.\n\nAnswer: Marta Hernández Romero",
    },
    {
        "id": "2hop__823584_776926",
        "context": '# Rotst\u00f6ckli\nThe Rotst\u00f6ckli (2,901 m) is a peak of the Urner Alps below the Titlis, on the border between the Swiss cantons of Obwalden and Nidwalden. It is Nidwalden\'s highest point. The summit is split between the municipalities of Engelberg (Obwalden) and Wolfenschiessen (Nidwalden).\n# Uri Alps\nThe Uri Alps (also known as "Urner Alps", ) are a mountain range in Central Switzerland and part of the Western Alps. They extend into the cantons of Obwalden, Valais, Bern, Uri and Nidwalden and are bordered by the Bernese Alps (Grimsel Pass) and the Emmental Alps to the west (the four lakes: Lungerersee, Sarnersee, Wichelsee, and Alpnachersee), the Schwyzer Alps to the north (Lake Lucerne), the Lepontine Alps to the south (the valley of Urseren with Andermatt) and the Glarus Alps to the east (Reuss).',
        "question": "What area contains the region that encompasses Rotst\u00f6ckli?",
        "generation": "Reasoning:\n- The context indicates that the Rotstöckli is a peak within the Urner Alps.\n- It further describes the Urner Alps as part of the Western Alps, a larger mountain range.\n- Therefore, the larger area that contains the region encompassing the Rotstöckli is the Western Alps, as deduced from the hierarchical geographical categorization provided.\n\nAnswer: Western Alps",
    },
]

answer_question_cot_fs = make_qa_func(
    system_prompt=SYSTEM_PROMPT_COT,
    few_shot_examples=FEW_SHOT_EXAMPLES_COT,
)

# %% ../../nbs/musique.qa.ipynb 17
SYSTEM_PROMPT_CTE = """
You are an excellent question-answering system known for providing accurate and reliable answers. Your responses should be solely based on the context information given, without drawing on prior knowledge.

Before answering the question, first, you extract relevant entity-relation-entity triplets from the context. Then, you answer the question based on the triplets.

# Output format
Triplets: [A list of entity-relation-entity triplets extracted from the context.]
Answer: [answer in least number of words possible]
""".strip()

answer_question_cte_zs = make_qa_func(
    system_prompt=SYSTEM_PROMPT_CTE,
)

# %% ../../nbs/musique.qa.ipynb 19
FEW_SHOT_EXAMPLES_CTE = [
    {
        "id": "2hop__784447_126070",
        "context": 'Glenhis Hern\u00e1ndez (born 7 October 1990 in Havana) is a taekwondo practitioner from Cuba. She was the 2013 World\nChampion in middleweight.\n\nThe current mayor of Havana ("President of the People\'s Power Provincial Assembly") is Marta Hern\u00e1ndez Romero, she\nwas elected on March 5, 2011.',
        "question": "Who is the current mayor of the city Glenhis Hern\u00e1ndez was born?",
        "generation": "Triplets: \nGlenhis Hern\u00e1ndez | birth place | Havana\nMarta Hern\u00e1ndez Romero | mayor of| Havana\n\nAnswer: Marta Hern\u00e1ndez Romero",
    },
    {
        "id": "2hop__823584_776926",
        "context": '# Rotst\u00f6ckli\nThe Rotst\u00f6ckli (2,901 m) is a peak of the Urner Alps below the Titlis, on the border between the Swiss cantons of Obwalden and Nidwalden. It is Nidwalden\'s highest point. The summit is split between the municipalities of Engelberg (Obwalden) and Wolfenschiessen (Nidwalden).\n# Uri Alps\nThe Uri Alps (also known as "Urner Alps", ) are a mountain range in Central Switzerland and part of the Western Alps. They extend into the cantons of Obwalden, Valais, Bern, Uri and Nidwalden and are bordered by the Bernese Alps (Grimsel Pass) and the Emmental Alps to the west (the four lakes: Lungerersee, Sarnersee, Wichelsee, and Alpnachersee), the Schwyzer Alps to the north (Lake Lucerne), the Lepontine Alps to the south (the valley of Urseren with Andermatt) and the Glarus Alps to the east (Reuss).',
        "question": "What area contains the region that encompasses Rotst\u00f6ckli?",
        "generation": "Triplets:\nRotst\u00f6ckli | part of | Urner Alps\nUrner Alps | part of | Western Alps\n\nAnswer: Western Alps",
    },
]

answer_question_cte_fs = make_qa_func(
    system_prompt=SYSTEM_PROMPT_CTE,
    few_shot_examples=FEW_SHOT_EXAMPLES_CTE,
)

# %% ../../nbs/musique.qa.ipynb 21
def load_qa_func(prompt_technique: str) -> Callable:
    prompt_technique = prompt_technique.lower()
    if prompt_technique == "standard":
        return answer_question_standard
    elif prompt_technique == "cot-zs":
        return answer_question_cot_zs
    elif prompt_technique == "cot-fs":
        return answer_question_cot_fs
    elif prompt_technique == "cte":
        return answer_question_cte_fs
    else:
        raise ValueError(f"Unknown prompt technique: {prompt_technique}")
