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

# %% auto 0
__all__ = ['log', 'DEFAULT_SYSTEM_PROMPT_TEMPLATE', 'DEFAULT_RELATION_SET_PROMPT_TEMPLATE', 'JERXChatFormatter']

# %% ../../nbs/jerx.prompt.ipynb 3
import random
from dataclasses import dataclass
from typing import Iterable, Callable, Any, Generator

from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/jerx.prompt.ipynb 6
DEFAULT_SYSTEM_PROMPT_TEMPLATE = """You are a helpful assistant that extracts up to {max_triplets} entity-relation-entity triplets from given text. Use '|' as delimiter and provide one triplet per line. The entities in a triplet must be different.
{relation_set_prompt}
""".strip()

DEFAULT_RELATION_SET_PROMPT_TEMPLATE = """Here are the list of relations that you can use:
{relation_set}
""".strip()

@dataclass
class JERXChatFormatter:
    system_prompt_template: str = DEFAULT_SYSTEM_PROMPT_TEMPLATE
    relation_set_prompt_template: str = DEFAULT_RELATION_SET_PROMPT_TEMPLATE
    relation_set: set|None = None
    few_shot_examples: list[dict]|None = None
    n_few_shot_examples: int = 3

    def __post_init__(self):
        if self.relation_set:
            self.relation_set = sorted(self.relation_set)
        if self.few_shot_examples is not None:
            self.few_shot_examples = list(self.few_shot_examples)

    def format(self, example: dict):
        few_shot_examples = self._choose_few_shot_examples()
        max_triplets = max(*[len(example['triplets']) for example in few_shot_examples], len(example['triplets'])) + 2
        messages = [
            {"role": "system", "content": self.make_system_message(max_triplets=max_triplets)},
            *list(self.make_messages(*few_shot_examples, example)),
        ]
        return {'conversations': messages}

    def make_system_message(self, max_triplets: int = 5) -> str:
        rsp = self.relation_set_prompt_template.format(relation_set=','.join(self.relation_set)) if self.relation_set else ""
        return self.system_prompt_template.format(max_triplets=max_triplets, relation_set_prompt=rsp)

    def make_messages(self, *examples) -> Generator[dict, None, None]:
        for example in examples:
            yield {"role": "user", "content": example["text"]}
            yield {"role": "assistant", "content": self._format_triplets(example["triplets"])}

    def _choose_few_shot_examples(self) -> list[dict]:
        if len(self.few_shot_examples) <= self.n_few_shot_examples:
            return self.few_shot_examples
        else:
            return random.sample(self.few_shot_examples, k=self.n_few_shot_examples)

    def _format_triplets(self, triplets: Iterable[str]) -> str:
        return '\n'.join(sorted(triplets))
