#!/bin/python

import sys, os
from pprint import pprint
from jinja2 import Environment, FileSystemLoader, meta
import yaml


sys.path.append(".")
os.chdir(os.path.dirname(os.path.abspath(__file__)))


if __name__ == "__main__":

    with open('prompt_config.yaml') as f:
        config = yaml.safe_load(f)

    print("prompt format:", config.get("prompt_format"))
    print(config)
    print()
    for prompt in config["prompts"]:
        if prompt["mode"] == "rag":
            print(f'--- prompt mode: {prompt["mode"]} ---')
            env = Environment(loader=FileSystemLoader("."))
            template = env.get_template(prompt["template"])

            source = template.environment.loader.get_source(template.environment, template.name)
            variables = meta.find_undeclared_variables(env.parse(source[0]))

            print("variables:", variables)
            print("---")

            data = {
                "query": "Quel est le meilleur moyen de cuire une blanquette?",
                "chunks" : [
                    {
                        "url": "http://data.gouv.fr",
                        "h": "hash49080805",
                        "title": "A chunk title",
                        "text": "Moi j'aime la blanquette avec du beurre dedans\nEt une sauce bien épaisse.",
                    },
                    {
                        "url": "http://...",
                        "h": "hash49080806",
                        "title": "A chunk title",
                        "text": "Faites chauffer la blanquette à feu doux pendant 46 heures.",
                        "context": "Recette de blanquette"
                    },
                                    {
                        "url": "http://...",
                        "h": "hash49080806",
                        "title": "A chunk title",
                        "text": "Les meilleures blanquettes se font avec des champignons de Paris",
                        "context": "Avis de grand-mère"
                    },
                                    {
                        "url": "http://...",
                        "h": "hash49080806",
                        "title": "A chunk title",
                        "text": """Étape 1 Faire revenir la viande dans un peu de beurre doux jusqu'à ce que les morceaux soient un peu dorés.
                                Étape 2: Saupoudrer de 2 cuillères de farine. Bien remuer.
                                Étape 3: Ajouter 2 ou 3 verres d'eau, les cubes de bouillon, le vin et remuer. Ajouter de l'eau si nécessaire pour couvrir.
                                Étape 4: Couper les carottes en rondelles et émincer les oignons puis les incorporer à la viande, ainsi que les champignons.
                                Étape 5: Laisser mijoter à feu très doux environ 1h30 à 2h00 en remuant.
                                Étape 6: Si nécessaire, ajouter de l'eau de temps en temps.
                                Étape 7: Dans un bol, bien mélanger la crème fraîche, le jaune d’oeuf et le jus de citron. Ajouter ce mélange au dernier moment, bien remuer et servir tout de suite.
                                """,
                        "context": "Recette Marmiton"
                    },

                ]
            }

            if "system_prompt" in variables:
                data["system_prompt"] = prompt["system_prompt"]
    
            rendered_template = template.render(**data)
            print(rendered_template)
            print("---")