# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/mhqa.llm.ipynb.

# %% auto 0
__all__ = ['log', 'DEFAULT_SYSTEM_PROMPT', 'USER_PROMPT', 'make_mhqa_chat', 'parse_llm_generation', 'QuestionAnsweringResult',
           'make_question_answer_func']

# %% ../../nbs/mhqa.llm.ipynb 4
import json
import os

from openai import OpenAI
from pydantic import BaseModel, Field

from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/mhqa.llm.ipynb 5
DEFAULT_SYSTEM_PROMPT = """You are an excellent Q&A system that is trusted around the world. You are given a question that requires multi-hop reasoning. Always answer the question using the provided context information, and not prior knowledge.

Some rules to follow:
1. Never directly reference the given context in your answer.
2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.

Output format:
Your output must be a single line in JSON such as:
{"reasoning": "Provide step by step multi-hop reasoning for the answer.", "answer": "Provide the final answer in 2-4 words."}
"""

USER_PROMPT = """The context information below is provided as a set of entity-relation-entity triplets from knowledge graph.
---------------------
{context}
---------------------
Given the context information and not prior knowledge, answer the question.
{question}
"""

def make_mhqa_chat(context: str, question: str) -> list[dict]:
    return [
            {
                "role": "system",
                "content": DEFAULT_SYSTEM_PROMPT,
            },
            {
                "role": "user",
                "content": USER_PROMPT.format(context=context, question=question),
            },
        ]

# %% ../../nbs/mhqa.llm.ipynb 6
def parse_llm_generation(output: str):
    return json.loads(output)

# %% ../../nbs/mhqa.llm.ipynb 7
class QuestionAnsweringResult(BaseModel):
    """Data model for answering the question."""

    reasoning: str = Field(description="Multi-hop reasoning for the answer.")
    answer: str = Field(description="The answer to the question in 2-4 words.")
    raw_output: str = Field(description="The raw output from the model.")


def make_question_answer_func(
    model_name: str = "gpt-3.5-turbo",
    client: OpenAI = None,
    completion_kwargs: dict | None = None,
):
    if client is None:
        client = OpenAI()

    if completion_kwargs is None:
        completion_kwargs = {}

    def func(context: str, question: str) -> QuestionAnsweringResult:
        messages = make_mhqa_chat(context, question)
        chat_completion = client.chat.completions.create(
            model=model_name,
            messages=messages,
            **completion_kwargs,
        )
        text = chat_completion.choices[0].message.content
        output = parse_llm_generation(text)
        return QuestionAnsweringResult(answer=output["answer"], reasoning=output["reasoning"], raw_output=text)

    return func
