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

# %% auto 0
__all__ = ['log', 'DEFAULT_SYSTEM_PROMPT', 'USER_PROMPT', 'QuestionAnsweringResult', 'make_question_answer_func',
           'RewardAssessment', 'make_reward_func']

# %% ../../../nbs/jerx.reward.llm.ipynb 3
import os
import instructor
from pydantic import BaseModel, Field
from openai import OpenAI

from ...text.utils import fuzzy_match
from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/jerx.reward.llm.ipynb 5
DEFAULT_SYSTEM_PROMPT = """You are an expert 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.
"""

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}
"""

class QuestionAnsweringResult(BaseModel):
    """Data model for answering the question."""

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


def make_question_answer_func(model_name: str = "gpt-3.5-turbo", client: OpenAI = None):
    if client is None:
        client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_API_BASE"))
    client = instructor.from_openai(client)

    def func(context: str, question: str) -> QuestionAnsweringResult:
        messages = [
            {
                "role": "system",
                "content": DEFAULT_SYSTEM_PROMPT,
            },
            {
                "role": "user",
                "content": USER_PROMPT.format(context=context, question=question),
            },
        ]
        return client.chat.completions.create(
            model=model_name,
            response_model=QuestionAnsweringResult,
            messages=messages,
        )

    return func

# %% ../../../nbs/jerx.reward.llm.ipynb 6
class RewardAssessment(QuestionAnsweringResult):
    reward: float = Field(description="The reward value for the answer.")

def make_reward_func(model_name: str = "gpt-3.5-turbo", answer_comparator=fuzzy_match):
    qa = make_question_answer_func(model_name)

    def reward(context: str, question: str, answers: list[str]) -> RewardAssessment:
        qa_result = qa(context, question)
        correct = any(answer_comparator(qa_result.answer, answer) for answer in answers)
        reward = 1.0 if correct else 0.0
        return RewardAssessment(**qa_result.dict(), reward=reward)

    return reward
