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

# %% auto 0
__all__ = ['log', 'DEFAULT_JERX_SYSTEM_MESSAGE_TEMPLATE', 'DEFAULT_JERX_USER_MESSAGE_TEMPLATE',
           'DEFAULT_JERX_ASSISTANT_MESSAGE_TEMPLATE', 'DEFAULT_JERX_CHAT_TEMPLATE', 'make_kg_triplet_extract_fn']

# %% ../../../nbs/jerx.fewshot.llm.ipynb 3
from llama_index.prompts import ChatPromptTemplate
from llama_index.llms import ChatMessage, OpenAI, LLM
from ..utils import parse_triplets
from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/jerx.fewshot.llm.ipynb 5
DEFAULT_JERX_SYSTEM_MESSAGE_TEMPLATE = """
Your task is to perform detailed entity-relation extraction from a document, creating a network of entities and their interrelations that enables answering complex, multi-hop questions. This requires careful analysis to identify and categorize entities (such as individuals, locations, organizations) and the specific, nuanced relationships between them.

Extract up to {max_knowledge_triplets} entity-relation-entity triplets from the given text.

# Guidelines
The goal is to build a detailed and accurate map of entities and their interrelations, enabling a comprehensive understanding of the document's content and supporting the answering of detailed, multi-hop questions derived from or related to the document. Prepare to adapt your extraction techniques to the nuances and specifics presented by the document, recognizing the diversity in structures and styles across documents.

# Core Objectives:
- **Comprehensive Entity and Relation Identification**: Systematically identify all relevant entities and their relationships within the document. Each entity and relation must be captured with precision, reflecting the document's depth of information.

- **Entity Differentiation and Categorization**: Distinguish between different types of entities, avoiding the amalgamation of distinct entities into a single category. For instance, separate individuals from their professions or titles and define their relationship clearly.

- **Clarification of Relationships and Avoidance of Redundancy**: Ensure each relationship is clearly defined, avoiding duplicate information. Relations should form a coherent, logical network, mapping connections between entities accurately, especially in hierarchical or geographic contexts.

- **Inference of Implicit Relations**: Infer and articulate relations that are implied but not explicitly stated within the document. This nuanced understanding allows for a richer, more interconnected entity-relation map.

- **Consistency and Cross-Validation**: Maintain consistency in entity references throughout the document and cross-validate entities and relations for accuracy. This includes harmonizing multiple references to the same entity and ensuring the entity-relation map is free from contradictions.

- **Detail-Oriented Relation Extraction**: Pay attention to the details within relations, capturing temporal and quantitative aspects where relevant. This adds depth to the understanding of each relationship, enhancing the capability to answer nuanced questions. Capture date and time relations with full detail as much as possible.

# Disambiguation and Unique Identification:
- **Explicit Disambiguation of Identical Names**: When encountering entities with identical names, explicitly disambiguate them by adding context-specific qualifiers in parentheses. These qualifiers should reflect the nature or category of the entity to prevent confusion and ensure clear differentiation. For example, differentiate geographical locations from non-geographical entities, people from non-person entities, and temporal from non-temporal entities with appropriate qualifiers.

# Formatting
- Avoid stopwords.
- Encode dates in the format: January 1, 1990
- Employ the format: `entity1 | relation | entity2` for each extracted relation, ensuring clarity and precision in representation. Each triplet should be in a new line.
""".strip()

DEFAULT_JERX_USER_MESSAGE_TEMPLATE = """
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()

DEFAULT_JERX_ASSISTANT_MESSAGE_TEMPLATE = """
Glenhis Hernández (Athlete) | born on | October 7, 1990
Glenhis Hernández (Athlete) | born in | Havana
Glenhis Hernández (Athlete) | specializes in | taekwondo
Glenhis Hernández (Athlete) | won | 2013 World Champion title (Middleweight)
Marta Hernández Romero (Politician) | serves as | mayor of Havana
Marta Hernández Romero (Politician) | holds | the position of "President of the People's Power Provincial Assembly"
Marta Hernández Romero (Politician) | elected on | March 5, 2011.
""".strip()

DEFAULT_JERX_CHAT_TEMPLATE = ChatPromptTemplate(
    message_templates=[
        ChatMessage(role='system', content=DEFAULT_JERX_SYSTEM_MESSAGE_TEMPLATE),
        ChatMessage(role='user', content=DEFAULT_JERX_USER_MESSAGE_TEMPLATE),
        ChatMessage(role='assistant', content=DEFAULT_JERX_ASSISTANT_MESSAGE_TEMPLATE),
        ChatMessage(role='user', content="{text}"),
    ]
)

# %% ../../../nbs/jerx.fewshot.llm.ipynb 6
def make_kg_triplet_extract_fn(
    *,
    max_knowledge_triplets: int = 10,
    llm: LLM | None = None,
):
    if llm is None:
        llm = OpenAI(model="gpt-3.5-turbo")

    def extract_kg_triplets(text: str) -> list[tuple[str, str, str]]:
        messages = DEFAULT_JERX_CHAT_TEMPLATE.format_messages(max_knowledge_triplets=max_knowledge_triplets, text=text)
        response = llm.chat(messages)
        return parse_triplets(response.message.content)

    return extract_kg_triplets
