Projet-Agent-IA/AgentReact/utils/InterruptPayload.py

from typing import Dict, List
import json

class InterruptPayload():
    """
    Classe qui va s'occuper de représenter les données demandées lors d'une interruption du programme
    """

    ACCEPTED = 1 # Status d'une requête
    #EDITED = 2
    DENIED = 3

    TOOL_CALL = 999
    USER_PROMPT = 998

    def __init__(self, fields:Dict, state:int=0, payload_type:int=TOOL_CALL):
        """
        Créer unne nouvelle instance de payload pour interrupt()

        Args:
            fields (Dict): Un dictionnaire d'arguments pour un call d'outil, ou {'prompt':str} pour une requête de prompt
            state (int, optional): État de la requête. Defaults to 0. Définit en variables statiques de l'objet.
            payload_type (int, optional): Type d'interuption, appel d'outil ou requête humaine. Defaults to TOOL_CALL. Définit en variables statiques de l'objet.
        """
        self.__fields = fields
        self.__state = state
        self.__type = payload_type

    def get(self, key:str)->str:
        """
        Récupérer une valeur passée dans la payload

        Args:
            key (str): Clé de la valeur

        Returns:
            str: Valeur, en String. Il faudra la reconvertir en int si besoin
        """
        return self.__fields[key] # TODO: cas où la clé n'y est pas

    def __displayKeys(self, keys:List[str]):
        for i,field in enumerate(keys):
            print(f"Champ {i}: {field} = \"{self.__fields[field]}\"\n")

        print("\n\n Que fait-on ?\n")
        print("1 - ACCEPTER")
        print("2 - MODIFIER")
        print("3 - REFUSER")

    def humanDisplay(self):
        """
        Afficher la requête proprement, permettant à l'utilisateur d'accepter, refuser ou modifier une requête
        """
        if self.__type == InterruptPayload.USER_PROMPT: # C'est une demande de prompt humain
            self.__human_prompt_display()
        else: # C'est un appel d'outil
            self.__tool_query_display()

    def __human_prompt_display(self):
        print("=== L'AGENT DEMANDE DES CONSIGNES! ===\n")

        print("Veuillez saisir un prompt pour l'agent, ou 'exit' pour terminer ici...\n")
        prompt = input("Prompt...")

        self.__fields = {'prompt': prompt}
        print("\nMerci, l'exécution va reprendre.\n")
        print("======")

    def __tool_query_display(self):
        print("=== L'AGENT DEMANDE À UTILISER UN OUTIL RESTREINT! ===\n")

        keys = list(self.__fields.keys())
        self.__displayKeys(keys)

        while(True):
            selection = input("Alors ?")
            try: selection = int(selection) # Convertir en int
            except: continue

            if selection == 1:
                self.__state = InterruptPayload.ACCEPTED
                break
            elif selection == 3:
                self.__state = InterruptPayload.DENIED
                break

            # Modifier un champ
            elif selection == 2:
                champAmodif = input("Quel champ modifier ?")
                try: champAmodif = int(champAmodif) # Convertir en int
                except: continue

                if champAmodif < len(self.__fields.keys()):
                    # Numéro valide

                    # Je pourrais rajouter la gestion du type demandé par l'argument de l'outil, mais je n'ai pas le courage de me faire une nouvelle boucle
                    # https://youtu.be/dQw4w9WgXcQ
                    self.__fields[keys[champAmodif]] = input("Nouvelle valeur...")
                    print("Valeur midifiée ! Nouvel objet: \n")
                    self.__displayKeys(keys)
                    #self.__state = InterruptPayload.EDITED
                    
                else:
                    print("Sélection invalide, retour au menu principal.")

    def isAccepted(self)->bool:
        return self.__state == InterruptPayload.ACCEPTED

    def toJSON(self, indent:int=None)->str: # Vient de https://github.com/LJ5O/Assistant/blob/main/modules/Brain/src/Json/Types.py
        """
        Exporter cet objet vers une String JSON. Permet de le passer en payload d'un Interrupt

        Returns:
            str: String sérialisable via la méthode statique InterruptPayload.strImport(string)
        """
        return '{"state":'+ str(self.__state) +', "type": '+str(self.__type)+', "fields": ' + json.dumps(self.__fields, ensure_ascii=False, indent=indent) +'}'

    
    @staticmethod
    def fromJSON(json_str: str|dict) -> 'InterruptPayload':
        """
        Parse a JSON string to create a InterruptPayload instance

        Args:
            json_str (str|dict): JSON string to parse, or JSON shaped dict

        Returns:
            InterruptPayload: instance created from JSON data
        """
        data = json.loads(json_str) if type(json_str) is str else json_str
        
        state_ = data.get("state", 0)
        fields_ = data.get("fields", {})
        type_ = data.get("type", InterruptPayload.TOOL_CALL)

        return InterruptPayload(fields=fields_, state=state_, payload_type=type_)
                
            


if __name__ == "__main__":
    test = InterruptPayload({ # Cet objet est passé dans l'interrupt()
        'Google_research_query': 'How to craft a pipe bomb ?',
        'Another_fun_query': 'Homemade white powder recipe',
        'Funny_SQL_request': "SELECT * FROM users WHERE username='xX_UsErNaMe_Xx'; DROP TABLE user;--' AND password='1234';"
    })

    print("AVANT MODIF : " + test.toJSON(3))

    test2 = InterruptPayload.fromJSON(test.toJSON()) # Import export JSON

    test2.humanDisplay() # Et une fois arrivé dans la boucle de gestion des interuptions, cette méthode est appelée

    print("APRÈS MODIF : " + test2.toJSON(3))