Chris4K/agent-reference-implemenation

1"""2Module: controller3 4This module provides a Controller class for handling user submissions and managing conversations.5 6Dependencies:7- app_agent_config: Module providing the AgentConfig class for configuring agents.8- utils.logger: Module providing logging functionalities.9- model.custom_agent: Module providing the CustomHfAgent class for interacting with Hugging Face models.10- model.conversation_chain_singleton: Module providing the ConversationChainSingleton class for managing conversation chains.11 12Classes:13- Controller: A class for handling user submissions and managing conversations.14"""15import os16from app_agent_config import AgentConfig  # Importing AgentConfig class from app_agent_config module17from utils.logger import log_response, IRCLogger     # Importing log_response function from utils.logger module18from model.custom_agent import CustomHfAgent  # Importing CustomHfAgent class from model.custom_agent module19from model.conversation_chain_singleton import ConversationChainSingleton  # Importing ConversationChainSingleton class from model.conversation_chain_singleton module20 21import logging22 23logging.basicConfig(level=logging.INFO)24 25server = "irc.efnet.org"26port = 666727nickname = "HFLogAAA"28channel = "#hflogs"29 30logger = IRCLogger(server, port, nickname, channel)31#logger.log_message("This is a test log message from the IRC logger.")32 33def cut_text_after_keyword(text, keyword):34    """35    Cuts text after the occurrence of a keyword.36 37    Args:38    - text (str): The text to be processed.39    - keyword (str): The keyword to search for in the text.40 41    Returns:42    - str: The processed text.43    """44    index = text.find(keyword)45    if index != -1:46        return text[:index].strip()47    return text48 49def get_text_after_last_occurrence(text, delimiter):50    """51    Retrieves the text after the last occurrence of the specified delimiter.52 53    Args:54    - text (str): The input text.55    - delimiter (str): The delimiter to search for.56 57    Returns:58    - str: The text after the last occurrence of the delimiter, or an empty string if the delimiter is not found.59    """60    last_index = text.rfind(delimiter)61    if last_index != -1:62        return text[last_index + len(delimiter):].strip()63    return ""64 65 66class Controller:67    """68    Controller class for handling user submissions and managing conversations.69    """70    def __init__(self):71        self.agent_config = AgentConfig()  # Initialize AgentConfig instance72       # logger.log_message("This is a test log message from the IRC logger.")73 74 75    image = []  # Class attribute for storing image data76 77    def handle_submission(self, user_message):78        """79        Handles user submission and interaction with the Hugging Face model.80 81        Args:82        - user_message (str): The message submitted by the user.83 84        Returns:85        - str: The response from the Hugging Face model.86        """87       # logger.log_message("This is a test log message from the IRC logger.")88        log_response("User input \n {}".format(user_message))89        log_response("selected_tools \n {}".format(self.agent_config.s_tool_checkboxes))90        log_response("url_endpoint \n {}".format(self.agent_config.url_endpoint))91        log_response("document \n {}".format(self.agent_config.document))92        log_response("image \n {}".format(self.agent_config.image))93        log_response("context \n {}".format(self.agent_config.context))94 95        selected_tools = [self.agent_config.tool_loader.tools[idx] for idx, checkbox in enumerate(self.agent_config.s_tool_checkboxes) if checkbox]96 97        agent = CustomHfAgent(98            url_endpoint=self.agent_config.url_endpoint,99            token=os.environ['HF_token'],100            additional_tools=selected_tools,101            input_params={"max_new_tokens": 192},102        )103 104        agent_response = agent.chat(user_message, document=self.agent_config.document, image=self.agent_config.image, context=self.agent_config.context)105 106        log_response("Agent Response\n {}".format(agent_response))107 108        return agent_response109 110    def handle_submission_chat(self, user_message, agent_response):111        """112        Handles user messages and responses in a conversation chain.113 114        Args:115        - user_message (str): The message submitted by the user.116        - agent_response (str): The response from the agent.117 118        Returns:119        - str: The response from the conversation chain.120        """121        agent_chat_bot = ConversationChainSingleton().conversation_chain("tmp")122        print(agent_chat_bot)123        print("------------ msg -----------------------")124        print(user_message + " ---- " )125        print("------------ /msg -----------------------")126        if agent_response is not None:127            msg = "[INST] You are a friendly chatbot who always responds to the user input in the style of a pirate. USER_INPUT:  "+user_message+" HINT: In a previous step the following was generated. use this to answer the user. AGENT_RESPONSE: "+ agent_response+" [/INST]"128            text = agent_chat_bot.predict(input=msg)129        else:130            msg = "[INST] You are a friendly chatbot who always responds to the user input in the style of a pirate. USER_INPUT:  "+user_message+"[/INST]"131            text = agent_chat_bot.predict(input=msg)132        print("----- msg----")133        print(msg)134        print("------------ text -----------------------")135        print(text)136        print("------------ /result -----------------------")137        result = get_text_after_last_occurrence(text, "AI: ")138        print(result)139        logger.log_message("Result: "+result+ " "+user_message)140        return result141