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

# %% auto 0
__all__ = ['log', 'DEFAULT_MODEL', 'DEFAULT_COMPLETION_KWARGS', 'FEW_SHOT_EXAMPLES', 'USER_PROMPT', 'SYSTEM_PROMPT_STANDARD',
           'SYSTEM_PROMPT_COT_FS', 'SYSTEM_PROMPT_CTE', 'answer_question_standard', 'answer_question_cot_fs',
           'answer_question_cot', 'answer_question_cte']

# %% ../../nbs/musique.qa.ipynb 4
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 6
FEW_SHOT_EXAMPLES = [
    {
        "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?",
        "answers": ["Marta Hern\u00e1ndez Romero"],
        "cte_generation": "Triplets: \nGlenhis Hern\u00e1ndez | birth place | Havana\nMarta Hern\u00e1ndez Romero | serves as | mayor of Havana\n\nAnswer: Marta Hern\u00e1ndez Romero",
        "cot_generation": "1. 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.\nMarta 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?",
        "cte_generation": "Triplets:\nRotst\u00f6ckli | part of | Urner Alps\nUrner Alps | part of | Western Alps\n\nAnswer: Western Alps",
        "cot_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",
    },
]

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

# %% ../../nbs/musique.qa.ipynb 10
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 2-4 words]
""".strip()

def answer_question_standard(
    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
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_STANDARD,
        },
        {
            "role": "user",
            "content": USER_PROMPT.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)

# %% ../../nbs/musique.qa.ipynb 13
SYSTEM_PROMPT_COT_FS = """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 2-4 words]
"""

def answer_question_cot_fs(
    context: str,
    question: str,
    examples: list[dict] = FEW_SHOT_EXAMPLES,
    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
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_COT_FS,
        },
    ]
    for example in examples:
        messages.append(
            {
                "role": "user",
                "content": USER_PROMPT.format(context=example["context"], question=example["question"]),
            }
        )
        messages.append({"role": "assistant", "content": example["cot_generation"]})

    messages.append(
        {
            "role": "user",
            "content": USER_PROMPT.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
    # Parse the response
    answer = ""
    reasoning = ""
    for line in generation.splitlines():
        if line.startswith("Answer:"):
            answer = line.split("Answer:")[1].strip()
        else:
            reasoning += line.replace("Reasoning:", "") + "\n"
    return dict(reasoning=reasoning.strip(), answer=answer, generation=generation)

# %% ../../nbs/musique.qa.ipynb 15
def answer_question_cot(
    context: str,
    question: str,
    model_name: str = DEFAULT_MODEL,
    completion_kwargs: dict | None = None,
    client=None,
) -> dict:
    return answer_question_cot_fs(context, question, [], model_name, completion_kwargs, client)

# %% ../../nbs/musique.qa.ipynb 18
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 2-4 words]
""".strip()

def answer_question_cte(
    context: str,
    question: str,
    examples: list[dict] = FEW_SHOT_EXAMPLES,
    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
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_CTE,
        },
    ]
    for example in examples:
        messages.append(
            {
                "role": "user",
                "content": USER_PROMPT.format(context=example["context"], question=example["question"]),
            }
        )
        messages.append(
            {
                "role": "assistant",
                "content": example["cte_generation"],
            }
        )
    messages.append(
        {
            "role": "user",
            "content": USER_PROMPT.format(context=context, question=question),
        },
    )
    
    # Generate the response
    chat_completion = client.chat.completions.create(
        model=model_name,
        messages=messages,
        **completion_kwargs,
    )
    generation = chat_completion.choices[0].message.content
    
    # Parse the response
    answer = ""
    triplets = []
    for line in generation.splitlines():
        if line.startswith("Answer:"):
            answer = line.split("Answer:")[1].strip()
        elif "|" in line:
            triplets.append(line.strip())
    return dict(triplets=triplets, answer=answer, generation=generation)
