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

# %% auto 0
__all__ = ['log', 'USER_PROMPT', 'EXAMPLE_CONTEXT', 'EXAMPLE_QUESTION', 'SYSTEM_PROMPT_STANDARD', 'SYSTEM_PROMPT_COT',
           'SYSTEM_PROMPT_COT_FS', 'EXAMPLE_COT_RESPONSE', 'SYSTEM_PROMPT_CTE', 'EXAMPLE_CTE_RESPONSE',
           'answer_question_standard', 'answer_question_cot', 'answer_question_cot_fs', 'answer_question_cte']

# %% ../../nbs/musique.qa.ipynb 4
import magentic
import openai
from pydantic import BaseModel, Field

from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/musique.qa.ipynb 6
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 7
EXAMPLE_CONTEXT = """
Glenhis Hernández (born 7 October 1990 in Havana) is a taekwondo practitioner from Cuba. She was the 2013 World
Champion in middleweight.

The current mayor of Havana ("President of the People's Power Provincial Assembly") is Marta Hernández Romero, she
was elected on March 5, 2011.
""".strip()

EXAMPLE_QUESTION = "Who is the current mayor of the city Glenhis Hernández was born?"


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

def answer_question_standard(
    context: str,
    question: str,
    model_name: str = "gpt-3.5-turbo",
    completion_kwargs: dict | None = None,
    client = None
) -> dict:
    
    if client is None:
        client = openai.Client()
    
    completion_kwargs = completion_kwargs or {}
    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
    if "Answer:" not in generation:
        return dict(answer="", generation=generation)
    answer = generation.split("Answer: ")[1].strip()
    return dict(answer=answer, generation=generation)

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

def answer_question_cot(
    context: str,
    question: str,
    model_name: str = "gpt-3.5-turbo",
    completion_kwargs: dict | None = None,
    client = None
) -> dict:

    if client is None:
        client = openai.Client()
    
    completion_kwargs = completion_kwargs or {}
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_COT,
        },
        {
            "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:", "").strip()
    return dict(reasoning=reasoning, answer=answer, generation=generation)

# %% ../../nbs/musique.qa.ipynb 14
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]
"""

EXAMPLE_COT_RESPONSE = """
Reasoning:
1. The context states that Glenhis Hernández was born in Havana.
2. It also mentions that the current mayor (President of the People's Power Provincial Assembly) of Havana is Marta Hernández Romero.
3. Marta Hernández Romero has been the mayor since her election on March 5, 2011.

Answer: Marta Hernández Romero
""".strip()

def answer_question_cot_fs(
    context: str,
    question: str,
    model_name: str = "gpt-3.5-turbo",
    completion_kwargs: dict | None = None,
    client = None
) -> dict:

    if client is None:
        client = openai.Client()
    
    completion_kwargs = completion_kwargs or {}
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_COT_FS,
        },
        {
            "role": "user",
            "content": USER_PROMPT.format(context=EXAMPLE_CONTEXT, question=EXAMPLE_QUESTION),
        },
        {
            "role": "assistant",
            "content": EXAMPLE_COT_RESPONSE
        },
        {
            "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 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 2-4 words]
"""

EXAMPLE_CTE_RESPONSE = """
Triplets: 
Glenhis Hernández | birth place | Havana
Marta Hernández Romero | serves as | mayor of Havana

Answer: Marta Hernández Romero
""".strip()

def answer_question_cte(
    context: str,
    question: str,
    model_name: str = "gpt-3.5-turbo",
    completion_kwargs: dict | None = None,
    client = None
) -> dict:

    if client is None:
        client = openai.Client()
    
    completion_kwargs = completion_kwargs or {}
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT_CTE,
        },
        {
            "role": "user",
            "content": USER_PROMPT.format(context=EXAMPLE_CONTEXT, question=EXAMPLE_QUESTION),
        },
        {
            "role": "assistant",
            "content": EXAMPLE_CTE_RESPONSE
        },
        {
            "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 = ""
    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)
