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

# %% auto 0
__all__ = ['DEFAULT_SYSTEM_PROMPT', 'USER_PROMPT', 'make_docs', 'make_qa_chat', 'parse_llm_generation', 'QuestionAnsweringResult',
           'make_question_answer_func', 'BaselineMultiHop', 'benchmark']

# %% ../../nbs/musique.multihopjerx.ipynb 3
import json
from typing import Callable

import pandas as pd
from tqdm.auto import tqdm

from ..jerx.reward.llm import QuestionAnsweringResult
from .eval import calculate_metrics, compare_answers

tqdm.pandas()

# %% ../../nbs/musique.multihopjerx.ipynb 4
def make_docs(example):
    ps = example["paragraphs"]
    for p in ps:
        idx = p["idx"]
        title = p["title"]
        body = p["paragraph_text"]
        is_supporting = p["is_supporting"]
        text = f"# {title}\n{body}"
        yield dict(
            text=text,
            is_supporting=is_supporting,
            parent_id=example["id"],
            idx=idx,
        )

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

import json
import os

from openai import OpenAI
from pydantic import BaseModel, Field

DEFAULT_SYSTEM_PROMPT = """You are an excellent Q&A system that is trusted around the world. 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 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_qa_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/qa.llm.ipynb 6
def parse_llm_generation(output: str):
    return json.loads(output)

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

    reasoning: str = Field(description="Concise 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_qa_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


# %% ../../nbs/musique.multihopjerx.ipynb 6
class BaselineMultiHop:
    def __init__(self, qa_func, retrieval_func):
        self.qa_func = qa_func
        self.retrieval_func = retrieval_func

    def _call(self, example) -> QuestionAnsweringResult:
        docs = list(make_docs(example))
        
        # First question
        question1 = example["question_decomposition"][0]["question"]
        query1 = question1
        docs1 = self.retrieval_func(docs, query1)
        context1 = "\n".join(doc['text'] for doc in docs1)
        result1 = self.qa_func(context=context1, question=question1)
        hop1 = {
            "question": question1,
            "query" : query1,
            "context": context1,
            "answer": result1.answer,
            "reasoning": result1.reasoning,
        }

        # Second question
        if result1.answer == "N/A":
            return QuestionAnsweringResult({
                "answer": "N/A",
                "reasoning": result1.reasoning,
                "hops": [hop1],
            })

        question2 = example["question_decomposition"][1]["question"]
        question2 = question2.replace("#1", result1.answer)
        query2 = question2
        docs2 = self.retrieval_func(docs, query2)
        context2 = "\n".join(doc['text'] for doc in docs2)
        result2 = self.qa_func(context=context2, question=question2)
        hop2 = {
            "question": question2,
            "query": query2,
            "context": context2,
            "answer": result2.answer,
            "reasoning": result2.reasoning,
        }
        return QuestionAnsweringResult(answer=result2.answer, reasoning=result2.reasoning, raw_output=json.dumps([hop1, hop2]))

    def __call__(self, example, ignore_errors: bool = False) -> QuestionAnsweringResult:
        try:
            output = self._call(example)
        except Exception as exc:
            if ignore_errors:
                id = example['id']
                print(f"Failed to answer the question {id}\n{exc}")
                output = QuestionAnsweringResult(reasoning="", answer="N/A", raw_output=str(exc))
            else:
                raise
        return output

# %% ../../nbs/musique.multihopjerx.ipynb 7
def benchmark(
    dataf: pd.DataFrame,
    qa_func: Callable,
    retrieval_func: Callable,
    ignore_errors: bool = False,
) -> tuple[pd.DataFrame, dict]:
    pipeline = BaselineMultiHop(qa_func, retrieval_func)

    def process(example):
        output = pipeline(example, ignore_errors=ignore_errors)
        example["predicted_answer"] = output.answer
        example["raw_llm_output"] = output
        return example

    dataf = dataf.progress_apply(process, axis=1)
    dataf = compare_answers(dataf)
    scores = calculate_metrics(dataf)
    scores["fuzzy_match"] = dataf["fuzzy_match"].mean()
    return dataf, scores
