Package glai

class AIMessage:
    """
    Represents a message in an AI system.

    Attributes:
        content (str): The content of the message.
        tag_open (str): The opening tag for the message.
        tag_close (str): The closing tag for the message.
    """

    def __init__(self, content:str, tag_open:str, tag_close:str):
        self.content = content
        self.tag_open = tag_open
        self.tag_close = tag_close

    def __str__(self) -> str:
        return f"{self.tag_open}{self.content}{self.tag_close}"
    
    def __repr__(self) -> str:
        return self.__str__()
    
    def edit(self, new_content, new_tag_open=None, new_tag_close=None):
        """
        Edits the content of the message.

        Args:
            new_content (str): The new content for the message.
        """
        self.content = new_content
        if new_tag_open is not None:
            self.tag_open = new_tag_open
        if new_tag_close is not None:
            self.tag_close = new_tag_close
    
    def text(self):
        """
        Returns the text representation of the message.

        Returns:
            str: The text representation of the message.
        """
        return self.__str__()
    
    def get_tags(self) -> tuple:
        """
        Returns the tags of the message.

        Returns:
            tuple: The tags of the message.
        """
        return (self.tag_open, self.tag_close)
    
    def to_dict(self) -> dict:
        """
        Returns the message as a dictionary.

        Returns:
            dict: The message as a dictionary.
        """
        return {
            "content": self.content,
            "tag_open": self.tag_open,
            "tag_close": self.tag_close,
        }
    
    @staticmethod 
    def from_dict(message_dict:dict) -> "AIMessage":
        """
        Creates a new AIMessage from a dictionary.

        Args:
            message_dict (dict): The dictionary to create the AIMessage from.

        Returns:
            AIMessage: The created AIMessage.
        """
        return AIMessage(message_dict["content"], message_dict["tag_open"], message_dict["tag_close"])

class AIMessages:
    """
    Represents a collection of messages in an AI system.
    Properties:
        user_tag_open (str): The opening tag for user messages.
        user_tag_close (str): The closing tag for user messages.
        ai_tag_open (str): The opening tag for AI messages.
        ai_tag_close (str): The closing tag for AI messages.
        system_tag_open (str): The opening tag for system messages.
        system_tag_close (str): The closing tag for system messages.
        messages (dict): The messages in the collection: {id: AIMessage}.
        _message_id_generator (int): The id generator for the messages.

    Args:
        messages (Union[AIMessages, AIMessage, str, list]): The messages to add to the collection.
        user_tags (tuple): The tags to use for user messages.
        ai_tags (tuple): The tags to use for AI messages.
        system_tags (tuple): The tags to use for system messages.
    """

    def __init__(self,user_tags:Union[tuple[str], list[str], dict]=("[INST]", "[/INST]"), ai_tags:Union[tuple[str], list[str], dict]=("", ""), system_tags:Optional[Union[tuple[str], list[str], dict]]=None):
        if isinstance(user_tags, dict):
            if "open" in user_tags and "close" in user_tags:
                self.user_tag_open = user_tags["open"]
                self.user_tag_close = user_tags["close"]
            else:
                raise ValueError(f"Invalid user tags: {user_tags}, for dict tags both 'open' and 'close' keys must be present.")
        elif isinstance(user_tags, set) or isinstance(user_tags, list) or isinstance(user_tags, tuple):
            self.user_tag_open = user_tags[0]
            self.user_tag_close = user_tags[1]
        else:
            raise TypeError(f"Invalid type for user tags: {type(user_tags)}, must be dict, set or list.")
        
        if isinstance(ai_tags, dict):
            if "open" in ai_tags and "close" in ai_tags:
                self.ai_tag_open = ai_tags["open"]
                self.ai_tag_close = ai_tags["close"]
            else:
                raise ValueError(f"Invalid user tags: {ai_tags}, for dict tags both 'open' and 'close' keys must be present.")
        elif isinstance(ai_tags, set) or isinstance(ai_tags, list) or isinstance(ai_tags, tuple):
            self.ai_tag_open = ai_tags[0]
            self.ai_tag_close = ai_tags[1]
        else:
            raise TypeError(f"Invalid type for user tags: {type(ai_tags)}, must be dict, set or list.")
        
        if system_tags is not None:
            if isinstance(system_tags, dict):
                if "open" in system_tags and "close" in system_tags:
                    self.system_tag_open = system_tags["open"]
                    self.system_tag_close = system_tags["close"]
                else:
                    raise ValueError(f"Invalid system tags: {system_tags}, for dict tags both 'open' and 'close' keys must be present.")
            elif isinstance(system_tags, set) or isinstance(system_tags, list) or isinstance(system_tags, tuple):
                self.system_tag_open = system_tags[0]
                self.system_tag_close = system_tags[1]
            else:
                raise TypeError(f"Invalid type for system tags: {type(system_tags)}, must be dict, set or list.")
        else:
            self.system_tag_open = None
            self.system_tag_close = None
        self.messages = {}
        self._message_id_generator = 0
    
    def user_tags(self) -> tuple[str]:
        """
        Returns the user tags.

        Returns:
            tuple[str]: The user tags.
        """
        return (self.user_tag_open, self.user_tag_close)
    
    def ai_tags(self) -> tuple[str]:
        """
        Returns the AI tags.

        Returns:
            tuple[str]: The AI tags.
        """
        return (self.ai_tag_open, self.ai_tag_close)

    def system_tags(self) -> Union[tuple[str], None]:
        """
        Returns the system tags.

        Returns:
            tuple[str]: The system tags.
        """
        if self.system_tag_open is not None and self.system_tag_close is not None:
            return (self.system_tag_open, self.system_tag_close)
        else:
            return None
    
    def load_messages(self, messages:Union[Any, AIMessage, str, list[Union[dict, AIMessage]]]) -> None:
        if messages is not None:
            if isinstance(messages, AIMessages):
                self.messages = messages.messages
            elif isinstance(messages, AIMessage):
                self.messages = [messages]
            elif isinstance(messages, str):
                self.messages = [AIMessage(messages, self.user_tag_open, self.user_tag_close)]
            elif isinstance(messages, list):
                if all([isinstance(message, AIMessage) for message in messages]):
                    self.messages = messages
                elif all(isinstance(message, str) for message in messages):
                    self.messages = [AIMessage(message, tag_open=self.user_tag_open, tag_close=self.user_tag_close) for message in messages]
                elif all(isinstance(message, dict) for message in messages):
                    self.messages = [AIMessage.from_dict(message_dict) for message_dict in messages]
                else:
                    raise TypeError("If passing list as messages it must be a list of AIMessage or str")
            else:
                raise TypeError("messages must be a list of AIMessage or str")
    
    def to_dict(self) -> dict:
        """
        Returns the messages as a dictionary.

        Returns:
            dict: The messages as a dictionary.
        """
        return {
            "user_tag_open": self.user_tag_open,
            "user_tag_close": self.user_tag_close,
            "ai_tag_open": self.ai_tag_open,
            "ai_tag_close": self.ai_tag_close,
            "system_tag_open": self.system_tag_open,
            "system_tag_close": self.system_tag_close,
            "messages": [message.to_dict() for message in self.messages]
        }
    
    @staticmethod
    def from_dict(messages_dict:dict) -> "AIMessages":
        """
        Creates a new AIMessages from a dictionary.

        Args:
            messages_dict (dict): The dictionary to create the AIMessages from.

        Returns:
            AIMessages: The created AIMessages.
        """
        ai_msgs = AIMessages()
        ai_msgs.user_tag_open = messages_dict["user_tag_open"]
        ai_msgs.user_tag_close = messages_dict["user_tag_close"]
        ai_msgs.ai_tag_open = messages_dict["ai_tag_open"]
        ai_msgs.ai_tag_close = messages_dict["ai_tag_close"]
        ai_msgs.system_tag_open = messages_dict["system_tag_open"] if "system_tag_open" in messages_dict else None
        ai_msgs.system_tag_close = messages_dict["system_tag_close"] if "system_tag_close" in messages_dict else None
        ai_msgs.load_messages(messages_dict["messages"])
        return ai_msgs
    
    def _generate_message_id(self) -> int:
        """
        Generates a unique message ID.
        Iters the message ID generator.

        Returns:
            int: The generated message ID.
        """
        self._message_id_generator += 1
        id = self._message_id_generator
        return id
    
    def add_message(self, message:Union[str, AIMessage], tag_open:str, tag_close:str) -> AIMessage:
        """
        Adds a new message to the messages dictionary.
        Iters the message ID generator.

        Parameters:
        - message (str|AIMessage): The content of the message or the message object
        - tag_open (str): The opening tag for the message, set to None if message is AIMessage
        - tag_close (str): The closing tag for the message, set to None if message is AIMessage

        Returns:
        AIMessage
        """
        if isinstance(message, str):
            message = AIMessage(message, tag_open, tag_close)
        self.messages[self._generate_message_id()] = message
        return self.messages[self._message_id_generator]

    def _insert_message(self, message:AIMessage, message_id:int) -> AIMessage:
        """
        Inserts a new message to the messages dictionary.
        Iters the message ID generator.

        Parameters:
        - message (AIMessage): The message object
        - message_id (int): The id to insert the message at

        Returns:
        AIMessage
        """
        self._message_id_generator += 1 if len(self.messages) > 0 else 0
        updated_messages = {}
        for id, msg in self.messages.items():
            if id < message_id:
                updated_messages[id] = msg
            else:
                updated_messages[id+1] = msg
        updated_messages[message_id] = message
        self.messages = updated_messages

    def add_user_message(self, message: Union[str, AIMessage]) -> AIMessage:
        """
        Adds a user message to the message list.
        Uses add_message() with the user tags.
        Automatically iters the message ID generator.

        Parameters:
            message (str): The message to be added.

        Returns:
            AIMessage
        """
        return self.add_message(message, self.user_tag_open, self.user_tag_close)

    def add_ai_message(self, message:Union[str, AIMessage]) -> AIMessage:
        """
        Adds a ai message to the message list.
        Uses add_message() with the ai tags.
        Automatically iters the message ID generator.

        Parameters:
            message (str): The message to be added.

        Returns:
            AIMessage
        """
        return self.add_message(message, self.ai_tag_open, self.ai_tag_close)
    
    def set_system_message(self, message:Union[str, AIMessage]) -> AIMessage:
        """
        Adds a system message to the message list.
        Uses add_message() with the system tags.
        Automatically iters the message ID generator.

        Parameters:
            message (str): The message to be added.

        Returns:
            AIMessage
        """
        if not self.has_system_tags():
            print("System tags are not set, this model does not support system messages.")
        else:
            if isinstance(message, str):
                message = AIMessage(message, self.system_tag_open, self.system_tag_close)    
            return self._insert_message(message, 0)

    def reset_messages(self) -> None:
        self.messages = {}
        self._message_id_generator = 0

    def __str__(self) -> str:
        return "".join([str(message) for message in self.messages.values()])

    def __repr__(self) -> str:
        return self.__str__()
    
    def text(self):
        return self.__str__()
    
    def get_last_message(self) -> AIMessage:
        """
        Returns the last message in the collection.

        Returns:
            AIMessage: The last message in the collection.
        """
        return self.messages[self._message_id_generator]
    
    def edit_last_message(self, new_content:str, tag_open:str=None, tag_close:str=None) -> None:
        """
        Updates the content of the last message in the collection.

        Parameters:
            new_content (str): The new content for the message.
            tag_open (str): Optional. The new opening tag for the message.
            tag_close (str): Optional. The new closing tag for the message.

        Returns:
            None
        """
        self.get_last_message().edit(new_content, tag_open, tag_close)
    
    def edit_message(self, message_id:int, new_content:str, tag_open:str=None, tag_close:str=None) -> None:
        """
        Updates the content of a message in the collection.

        Parameters:
            message_id (int): The id of the message to edit.
            new_content (str): The new content for the message.
            tag_open (str): The new opening tag for the message.
            tag_close (str): The new closing tag for the message.

        Returns:
            None
        """
        self.messages[message_id].edit(new_content, tag_open, tag_close)

    def edit_system_message(self, new_content:str) -> None:
        if self.system_tags() is None:
            raise ValueError("System tags are not set, this model does not support system messages.")
        else:
            if self.messages[0].tag_open == self.system_tag_open and self.messages[0].tag_close == self.system_tag_close:
                self.edit_message(0, new_content)
            else:
                print("Warning: System message not found, adding system message to the start of the message list.")
                self.set_system_message(new_content)

    def has_system_tags(self) -> bool:
        """
        Returns whether the model supports system messages.

        Returns:
            bool: Whether the model supports system messages.
        """
        return self.system_tags() is not None
    
    def save_json(self, file_path:str) -> None:
        """
        Saves the messages as a json file.

        Parameters:
            file_path (str): The path to save the json file to.

        Returns:
            None
        """
        save_json_file(self.to_dict(), file_path)
    
    @staticmethod
    def from_json(file_path:str) -> "AIMessages":
        """
        Creates a new AIMessages from a json file.

        Args:
            file_path (str): The path to the json file.

        Returns:
            AIMessages: The created AIMessages.
        """
        return AIMessages.from_dict(load_json_file(file_path))
    
    @staticmethod
    def create_single_message(message:str, tag_open:str="", tag_close:str="") ->AIMessage:
        """
        Creates a single AIMessage.

        Args:
            message (str): The content of the message.
            tag_open (str): The opening tag for the message.
            tag_close (str): The closing tag for the message.

        Returns:
            AIMessage: The created AIMessage.
        """
        return AIMessage(message, tag_open, tag_close)

class AutoAI:
    """
    Initialize the AutoAI class for super easy LLM AI generation.

    Searches for a model based on name/quantization/keyword. 
    Downloads model and sets up LlamaAI.

    Args:
        name_search: Name of model to search for. Optional. Default None.
        quantization_search: Quantization of model to search for. Optional. Default None.
        keyword_search: Keyword of model to search for. Optional. Default None.
        new_tokens: New token length for LlamaAI model. Default 1500.
        max_input_tokens: Max input tokens for LlamaAI model. Default 900.
        model_db_dir: Directory to store model data in. Defaults to global packages model directory.

    Attributes:
        model_db: ModelDB object. - represents the database of models, has useful functions for searching and importing models.
        model_data: ModelData object. - represents the data of the model, has useful functions for creating, downloading and loading the model data and gguf.
        ai: LlamaAI object. - represents the LlamaAI model, a wrapper for llama llm and tokenizer models quantized to gguf format. Has methods for adjusting generation and for generating.
        msgs: AIMessages object. - represents the AIMessages a collection of AIMessage objects, has useful functions for adding and editing messages and can be printed to string.
        
    """
    def __init__(self, 
                 name_search: Optional[str] = None,
                 quantization_search: Optional[str] = None,
                 keyword_search: Optional[str] = None,
                 search_only_downloaded_models:bool = False,
                 max_total_tokens: int = 1500,
                 model_db_dir:Optional[str] = None,
                 ) -> None:

        self.model_db = ModelDB(model_db_dir=model_db_dir, copy_verified_models=True)
        self.model_data: ModelData = self.model_db.find_model(
            name_search, quantization_search, keyword_search, search_only_downloaded_models
        )
        self.model_data.download_gguf()
        self.ai = LlamaAI(
            self.model_data.gguf_file_path, max_tokens=max_total_tokens
        )
        print(f"Using model: {self.model_data}")
        self.msgs: AIMessages = AIMessages(
            self.model_data.user_tags, self.model_data.ai_tags, self.model_data.system_tags
        )
    
    def generate_from_messages(self, stop_at:str = None, include_stop_str:bool = True) -> AIMessage:
        prompt = self.msgs.text()
        ai_message = self.generate_from_literal_string(prompt, stop_at=stop_at, include_stop_str=include_stop_str)
        self.msgs.add_ai_message(ai_message)
        return ai_message
    
    def generate_from_literal_string(
        self, 
        prompt: str,
        stop_at:str = None,
        include_stop_str:bool = True
    ) -> AIMessage:
        """
        Generate text from a prompt using the LlamaAI model.

        Args:
            prompt: Prompt text to generate from.

        Returns:
            Generated text string.
        """
        return self.ai.infer(
            prompt, only_string=True, stop_at_str=stop_at, include_stop_str=include_stop_str
        )

    def generate(
        self,
        user_message: str,
        ai_message_tbc: Optional[str] = None,
        stop_at:Optional[str] = None,
        include_stop_str:bool = True,
        system_message: Optional[str] = None
    ) -> AIMessage:
        """
        Generate an AI response to a user message.

        Args:
            user_message: User message text.
            ai_message_tbc: Optional text to prepend.
            stop_at: Optional string to stop generation at.
            include_stop_str: Whether to include the stop string in the generated message.
            system_message: Optional system message to include at the start, not all models support this.
            If you provide system message to a model that doesn't support it, it will be ignored.
            You can check if a model supports system messages by checking the model_data.has_system_tags() method.
        Returns:
            Generated AIMessage object.
        """
        generation_messages = AIMessages(user_tags=self.model_data.user_tags, ai_tags=self.model_data.ai_tags, system_tags=self.model_data.system_tags)
        generation_messages.reset_messages()
        if self.model_data.has_system_tags():
            if system_message is not None:
                generation_messages.set_system_message(system_message)
            else:
                print("WARNING: Model seeps to support system messages, but no system message provided.")
        generation_messages.add_user_message(user_message)


        if ai_message_tbc is not None:
            generation_messages.add_message(
                ai_message_tbc, 
                self.msgs.ai_tag_open, 
                ""
            )
        print(f"Promt: {generation_messages.text()}")
        
        generated = self.generate_from_literal_string(generation_messages.text(), stop_at=stop_at, include_stop_str=include_stop_str)

        if ai_message_tbc is not None:
            generation_messages.edit_last_message(
                ai_message_tbc + generated,
                self.msgs.ai_tag_open,
                self.msgs.ai_tag_close
            )
        else:
            generation_messages.add_ai_message(generated)

        print(f"Generated: {generation_messages.get_last_message().text()}")
        output = generation_messages.get_last_message()
        return output

    def count_tokens(
        self,
        user_message: str,
        ai_message_tbc: Optional[str] = None
    ) -> int:
        """
        Count the number of tokens in a generated message.

        Args:
            user_message: User message text.
            ai_message_tbc: Optional text to prepend.

        Returns:
            Number of tokens in generated message.
        """
        generation_messages = AIMessages()
        generation_messages.reset_messages()
        generation_messages.add_user_message(user_message)

        if ai_message_tbc is not None:
            generation_messages.add_message(
                ai_message_tbc, 
                self.msgs.ai_tag_open, 
                ""
            )
        return self.ai.count_tokens(generation_messages.text())
    
    def is_within_input_limit(
        self,
        user_message: str,
        ai_message_tbc: Optional[str] = None
        ) -> bool:
        """
        Check if the generated message is within the input limit.

        Args:
            user_message: User message text.
            ai_message_tbc: Optional text to prepend.

        Returns:
            True if within input limit, False otherwise.
        """
        return self.ai.is_within_input_limit(self.count_tokens(user_message, ai_message_tbc))

class EasyAI:
    """
    EasyAI provides a simple interface for LlamaAI based AI using quantized GGUF models
    for inference on CPU.
    
    Initialization:
    Can be intialized with no arguments, followed by following configuration methods, step by step, all in one or using a dict:\n
    Step by step call these with appropriate arguments:\n
        1. `self.load_model_db(model_db_dir: str = MODEL_EXAMPLES_DB_DIR)`
        2. One of the following with necessary args: `self.model_data_from_url()` or `self.model_data_from_file()` or `self.find_model_data()`
        3. `self.load_ai(max_total_tokens: int = 200)`
    All in one:\n
        ```python
        self.configure(
            model_db_dir:str = MODEL_EXAMPLES_DB_DIR, 
            model_url: Optional[str] = None, 
            model_gguf_path: Optional[str] = None, 
            name_search: Optional[str] = None, 
            quantization_search: Optional[str] = None, 
            keyword_search: Optional[str] = None, 
            max_total_tokens: int = 200, 
            )
        ```
        
        Or with a dictionary with the following keys:
        - model_db_dir: Directory to store model data in. Defaults to MODEL_EXAMPLES_DB_DIR.
        - model_url: URL of model to configure with. Automatically downloads and builds as needed. (Optional)
        - name_search: Name of model to search for in the model db dir.(Optional)
        - quantization_search: Quantization of model to search for in the model db dir..(Optional)
        - keyword_search: Keyword of model to search for in the model db dir..(Optional)
        - model_gguf_path: Path to GGUF file of model to configure with.(Optional, not a recommended method, doesn't preserve download url)
        - max_total_tokens: Max tokens to be processed by LlamaAI model. (Defaults to 200, set to around 500-1k for regular use)

    Attributes:
        model_db: ModelDB for searching/loading models
        messages: AIMessages for tracking conversation 
        model_data: ModelData of selected model
        lai: LlamaAI instance for generating text

    Methods:
        DB:
            load_model_db: Load ModelDB from directory
        ModelData:
            find_model_data: Search model DB for ModelData
            model_data_from_url: Get ModelData from URL
            model_data_from_file: Load ModelData from file
        Load to memory:
            load_ai: Create LlamaAI instance from ModelData
        Inference:
            infer: Generate AI response to user message

    EasyAI handles loading models, setting up messages/LLamaAI,
    and generating responses. It provides a simple interface to using
    LLama
    """
    def __init__(self, **kwds) -> None:
        self.model_db: ModelDB = None
        self.messages: Optional[AIMessages] = None
        self.model_data: Optional[ModelData] = None
        self.ai: Optional[LlamaAI] = None
        if kwds:
            self.configure(**kwds)

    def configure(self,
                  model_db_dir: Optional[str] = None,
                  model_url: Optional[str] = None,
                  model_gguf_path: Optional[str] = None,
                  name_search: Optional[str] = None,
                  quantization_search: Optional[str] = None,
                  keyword_search: Optional[str] = None,
                  search_only_downloaded: bool = False,
                  max_total_tokens: int = 200,
                                            ) -> None:
        """
        Configure EasyAI with model data.

        Configures EasyAI with model data from given URL, GGUF file path, or model name/quantization/keyword.
        Sets model data attribute.

        Args:
            model_db_dir: Directory to store model data in. If none is provided global db is used.This is preferred for most use cases.
            max_total_tokens: Max tokens to be processed (input+generation) by LlamaAI model. (Defaults to 200, set to around 500-1k for regular use)
            
            Provide at least one of these args to fetch ModelData: 
            ---
            model_url: URL of model to configure with. Automatically downloads and builds as needed. (Optional)
            name_search: Name of model to search for in the model db dir.(Optional)
            quantization_search: Quantization of model to search for in the model db dir..(Optional)
            keyword_search: Keyword of model to search for in the model db dir..(Optional)
            model_gguf_path: Path to GGUF file of model to configure with.(Optional, not a recommended method, doesn't preserve download url)
            ---

        Raises:
            Exception: If no model DB loaded.
            Exception: If no model data found.

        """
        if model_db_dir is None:
            print(f"Using provided verified models DB, files at {VERIFIED_MODELS_DB_DIR}")
        self.load_model_db(model_db_dir)
        if model_url is not None:
            self.model_data_from_url(model_url)
        elif model_gguf_path is not None:
            self.model_data_from_file(model_gguf_path)
        elif name_search is not None or quantization_search is not None or keyword_search is not None:
            self.find_model_data(name_search, quantization_search, keyword_search, search_only_downloaded)
        else:
            raise Exception("Can't find model data. Please provide a model URL, GGUF file path, or model name/quantization/keyword.")
        
        self.load_ai(max_total_tokens)
    


            
            


    def load_model_db(self, db_dir:Optional[str] = None, copy_verified_models=True) -> None:
        """
        Load ModelDB from given directory.

        Args:
            db_dir: Directory to load ModelDB from.
            copy_examples: Whether to copy example GGUF files to db_dir if db_dir is empty.
        """
        self.model_db = ModelDB(model_db_dir=db_dir, copy_verified_models=copy_verified_models)

    def import_verified_models_to_db(self, model_name_quantization_list:Optional[list[Union[list,set,tuple]]] = None) -> None:
        """
        Import verified models to new ModelDB.

        Imports verified models from global verified models DB to new ModelDB. If model_name_quantization_list is provided, only models in the list are imported.

        Args:
            model_name_quantization_list: List of tuples of model name and quantization to import. If None, all models are imported.
        """
        if self.model_db is None:
            raise Exception("No model DB loaded. Use load_model_db() first.")
        if model_name_quantization_list is not None:
            for n_q_data in model_name_quantization_list:
                name = n_q_data[0]
                quantization = n_q_data[1]
                self.model_db.import_verified_model(name, quantization, None, True)
        else:
            self.model_db.import_verified_model(None, None, None, True)
        
    def find_model_data(self,
                        model_name: Optional[str] = None,  
                        quantization: Optional[str] = None,
                        keyword: Optional[str] = None,
                        only_downloaded: bool = False) -> ModelData:
        """
        Find model data in database that matches given criteria.

        Searches model database for model data matching the given model name, 
        quantization, and/or keyword. Any parameters left as None are not used  
        in the search.

        Args:
            model_name: Name of model to search for.
            quantization: Quantization of model to search for.
            keyword: Keyword of model to search for.

        Returns:
            ModelData object if a match is found, else None.
        """
        if self.model_db is None:
            raise Exception("No model DB loaded. Use load_model_db() first.")
        model_data = self.model_db.find_model(model_name, quantization, keyword, only_downloaded)
        self.model_data = model_data
        return model_data

    def model_data_from_url(self,
                            url: str,
                            user_tags: Tuple[str, str] = ("", ""),
                            ai_tags: Tuple[str, str] = ("", ""),
                            description: Optional[str] = None,
                            keywords: Optional[str] = None,
                            save: bool = True) -> None:
        """
        Get model data for URL, downloading model if needed.

        Checks if model data already exists for the given URL. If not, downloads
        the model from the URL and creates new model data. Sets model data attribute.

        Args:
            url: URL of model to get data for.
            user_tags: User tags to assign if creating new model data. 
            ai_tags: AI tags to assign if creating new model data.
            description: Optional description for new model data.
            keyword: Optional keyword for new model data.
            save: Whether to save new model data JSON file.
        """
        if self.model_db is None:
            raise Exception("No model DB loaded. Use load_model_db() first.")
        print(f"Trying to get model data from url: {url}")
        print(f"Checking if model data already exists...")
        model_data = self.model_db.get_model_by_url(url)
        if model_data is None:
            print(f"Model data not found. Creating new model data...")
            model_data = ModelData(url, self.model_db.gguf_db_dir, user_tags, ai_tags, description, keywords)
            print(f"Created model data: {model_data}")
        else:
            print(f"Found model data: {model_data}")
        if save:
            model_data.save_json()
        self.model_data = model_data

    def model_data_from_file(self,
                             gguf_file_path: str,
                             user_tags: Tuple[str, str] = ("", ""),
                             ai_tags: Tuple[str, str] = ("", ""),
                             description: Optional[str] = None,
                             keyword: Optional[str] = None,
                             save: bool = False) -> None:
        """
        Get model data from local GGUF file.

        Loads model data from the given local GGUF file path. Sets model data attribute.

        Args:
            gguf_file_path: Path to GGUF file.
            user_tags: User tags to assign to model data.
            ai_tags: AI tags to assign to model data.
            description: Optional description for model data.
            keyword: Optional keyword for model data.
            save: Whether to save model data JSON file.
        """
        if self.model_db is None:
            raise Exception("No model DB loaded. Use load_model_db() first.")
        model_data = ModelData.from_file(gguf_file_path, self.model_db.gguf_db_dir, user_tags, ai_tags, description, keyword)
        if save:
            model_data.save_json()
        self.model_data = model_data

    def _load_messages(self) -> None:
        """
        Load AIMessages using tags from model data.

        Uses user_tags and ai_tags from loaded model data to initialize AIMessages object. 
        Sets messages attribute.

        Raises:
            Exception: If no model data loaded yet.
        """
        if self.model_data is None:
            raise Exception("No model data loaded. Use find_model_data(), get_model_data_from_url(), or get_model_data_from_file() first.")
        self.messages = AIMessages(user_tags=self.model_data.user_tags, ai_tags=self.model_data.ai_tags, system_tags=self.model_data.system_tags)

    def load_ai(self,
                max_total_tokens: int = 200,) -> None:
        """
        Load LlamaAI model from model data.

        Downloads model file from model data URL if needed. Initializes LlamaAI with model and sets lai attribute.

        Args:
            max_total_tokens: Max tokens for LlamaAI model.
        Raises:
            Exception: If no model data or messages loaded yet.
        """
        self._load_messages()
        if self.messages is None:
            raise Exception("No messages loaded. Use load_messages() first.")
        if self.model_data is None:
            raise Exception("No model data loaded. Use find_model_data(), get_model_data_from_url(), or get_model_data_from_file() first.")
        self.model_data.download_gguf()
        self.ai = LlamaAI(self.model_data.model_path(), max_tokens=max_total_tokens)
        print(f"Loaded: {self.model_data}")

    def generate(self,
              user_message: str,
              ai_message_tbc: Optional[str] = None,
              stop_at:Optional[str]=None,
              include_stop_str:bool=True,
              system_message: Optional[str] = None
              ) -> AIMessage:
        """
        Generate AI response to user message.

        Runs user message through loaded LlamaAI to generate response. Allows prepending optional 
        content to AI response. Adds messages and returns generated AIMessage.

        Args:
            user_message: User message text.
            ai_message_tbc: Optional text to prepend to AI response.
            stop_at: Optional string to stop generation at.
            include_stop_str: Whether to include stop string in generated message.
            system_message: Optional system message to include at the start, not all models support this.
            If you provide system message to a model that doesn't support it, it will be ignored.
            You can check if a model supports system messages by checking the model_data.has_system_messages()

        Returns:
            Generated AIMessage object.

        Raises:
            Exception: If no AI or messages loaded yet.
        """
        if self.ai is None:
            raise Exception("No AI loaded. Use load_ai() first.")
        if self.messages is None:
            raise Exception("No messages loaded. Use load_ai() first.")
        self.messages.reset_messages()
        if self.model_data.has_system_tags():
            if system_message is not None:
                self.messages.set_system_message(system_message)
            else:
                print("WARNING: Model supports system messages, but no system message provided.")
        self.messages.add_user_message(user_message)
        print(f"Input to model: \n{self.messages.text()}")
        generated: str = ""
        if ai_message_tbc is not None:
            generated += ai_message_tbc
            self.messages.add_message(ai_message_tbc, self.model_data.get_ai_tag_open(), "")
        if stop_at is None:
            stop_at = self.messages.ai_tag_close if any([self.messages.ai_tag_close is None, self.messages.ai_tag_close == "", self.messages.ai_tag_close != " "]) else None
            include_stop_str = False
        generated += self.ai.infer(self.messages.text(), only_string=True, stop_at_str=stop_at, include_stop_str=include_stop_str)
        if ai_message_tbc is not None:
            self.messages.edit_last_message(generated,
                                            self.model_data.get_ai_tag_open(),
                                            self.model_data.get_ai_tag_close())
        else:
            self.messages.add_ai_message(generated)
        print(f"AI message: \n{self.messages.get_last_message()}")
        return self.messages.get_last_message()

    



    def count_tokens(
        self,
        user_message_text: str,
        ai_message_tbc: Optional[str] = None
    ) -> int:
        """
        Count the number of tokens in a generated message.

        Args:
            user_message_text: User message text.
            ai_message_tbc: Optional text to prepend.

        Returns:
            Number of tokens in generated message.
        """
        generation_messages = AIMessages()
        generation_messages.reset_messages()
        generation_messages.add_user_message(user_message_text)

        if ai_message_tbc is not None:
            generation_messages.add_message(
                ai_message_tbc, 
                self.messages.ai_tag_open, 
                ""
            )

        return self.ai.count_tokens(generation_messages.text())
    
    def is_within_context(self,
        prompt: str,
        ) -> bool:
        """
        Check if the generated message is within the input limit.

        Args:
            user_message_text: User message text.
            ai_message_tbc: Optional text to prepend.

        Returns:
            True if within input limit, False otherwise.
        """
        return self.ai.is_prompt_within_limit(prompt)
    
    def import_from_repo(self, hf_repo_url: str, user_tags: Tuple[str, str] = ("", ""), ai_tags: Tuple[str, str] = ("", ""), system_tags: Optional[Tuple[str, str]] = (None, None), keywords: Optional[str] = None, description: Optional[str] = None, replace_existing: bool = False) -> None:
        """
        Imports model data from HuggingFace model repo to current model DB. 

        Args:
            hf_repo_url: URL of model to import.
            user_tags: User tags to assign to model data.
            ai_tags: AI tags to assign to model data.
            system_tags: System tags to assign to model data.
            description: Optional description for model data.
            keyword: Optional keyword for model data.
            replace_existing: Whether to replace existing model data if found.

        Raises:
            Exception: If no model DB loaded yet.
        """



        self.model_db.import_models_from_repo(
            hf_repo_url=hf_repo_url,
            user_tags=user_tags,
            ai_tags=ai_tags,
            system_tags=system_tags,
            keywords=keywords,
            description=description,
            replace_existing=replace_existing,
        )