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

# %% auto 0
__all__ = ['log', 'SYSTEM_PROMPT', 'Assessment', 'assess_quality']

# %% ../../../nbs/jerx.reward.llm.ipynb 3
import magentic
from pydantic import BaseModel, Field

from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/jerx.reward.llm.ipynb 5
SYSTEM_PROMPT = """
Assess the quality of the entity-relation-entity triplets extracted from the document. Provide a brief explanation for your assessment and a score between 0 and 1. The score should reflect the quality of the triplets in terms of their relevance and completeness. The explanation should describe the strengths and weaknesses of the triplets and how they could be improved. Your output must 
Reasoning: 1-line explanation of the assessment.
Score: a number between 0 and 1.
""".strip()

@magentic.chatprompt(
        magentic.SystemMessage(SYSTEM_PROMPT),
        magentic.UserMessage("""
# Document
Washington Township is one of the thirteen townships of Henry County, Ohio, United States. As of the 2010
census the population was 1,912, of whom 1,794 lived in the unincorporated portion of the township.
# Triplets
Washington Township | is one of | thirteen townships of Henry County
Washington Township | located in | Henry County, Ohio
Washington Township | had population in 2010 census | 1,912
Unincorporated portion of Washington Township | had population in 2010 census | 1,794
Washington Township | part of census | 2010
""".strip()),
        magentic.AssistantMessage("""
Reasoning: The triplets are well-crafted and highly relevant but miss some opportunities for providing comprehensive details that would make them even more useful for answering a broader range of related questions.
Score: 0.6
""".strip()),
        magentic.UserMessage("""
# Document
Turkeycock Mountain is a mountain summit located in Franklin County, Virginia and Henry County, Virginia.
Rising out of the eastern foothills of the Blue Ridge Mountains, Turkeycock Mountain rises to above sea level and is
located at . A portion of the mountain is protected as the Turkeycock Wildlife Management Area, which is open to the
public.

# Triplets
Turkeycock Mountain | located in | Franklin County, Virginia
Turkeycock Mountain | located in | Henry County, Virginia
Turkeycock Mountain | part of | Blue Ridge Mountains
Turkeycock Mountain | feature | rises to above sea level
Turkeycock Mountain | rises from | eastern foothills of Blue Ridge Mountains
Turkeycock Mountain | elevation | above sea level
Turkeycock Wildlife Management Area | part of | Turkeycock Mountain
Turkeycock Wildlife Management Area | status | protected
Turkeycock Wildlife Management Area | open to | public
Turkeycock Wildlife Management Area | allows | public access
""".strip()),
        magentic.AssistantMessage("""
Reasoning: The triplets are accurate and complete.
Score: 1.0
""".strip()),
        magentic.UserMessage("""
# Document
Joachim II ( or "Hektor"; 13 January 1505 – 3 January 1571) was a Prince-elector of the Margraviate of Brandenburg (1535–1571), the sixth member of the House of Hohenzollern. Joachim II was the eldest son of Joachim I Nestor, Elector of Brandenburg and his wife Elizabeth of Denmark, Norway, and Sweden. He received the cognomen Hector after the Trojan prince and warrior for his qualities and prowess.

# Triplets
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in Europe
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in Europe
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in Europe
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in Europe
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
Joachim II | legacy | helped to establish the Hohenzollern dynasty as a major power in the world
""".strip()),
        magentic.AssistantMessage("""
Reasoning: There are many duplicate triplets in the response and the triplets do not capture the full range of information in the document. The format of each triplet is correct.
Score: 0.05
""".strip()),

        magentic.UserMessage("""
# Document
{document}

# Triplets
{triplets_str}
""".strip()),
)
def _assess_triplets(document: str, triplets_str: str) -> str: ...


# %% ../../../nbs/jerx.reward.llm.ipynb 6
class Assessment(BaseModel):
    reasoning: str = Field(description="A brief reasoning for the assessment.")
    reward: float = Field(ge=0, le=1, description="The reward of the assessment.")

def assess_quality(document: str, triplets_str: str) -> Assessment:
    result = _assess_triplets(document, triplets_str)
    lines = result.splitlines()
    if len(lines) < 2:
        return Assessment(reasoning=result, reward=0.0)
    
    last_line = lines[-1]
    if not last_line.startswith("Score: "):
        return Assessment(reasoning=result, reward=0.0)
    
    reward = float(last_line[len("Score:"):].strip())
    return Assessment(reasoning="\n".join(lines[:-1]), reward=reward)
