"""
This module defines a custom LLM class compatible with LangChain
that wraps the `Agent` class. This class allows BRAD to be integrated into 
LangChain's workflow and used as a standard LLM. The class overrides core methods required 
by LangChain to interact with BRAD, including handling prompts, stopping criteria, and 
managing run callbacks.

Classes:
    BradLLM: A subclass of LangChain's LLM that wraps BRAD to provide a LangChain-compatible interface.

Attributes:
    bot (Optional[Any]): An instance of the BRAD chatbot. This attribute is expected to be set externally
        and will handle the actual prompt invocation.

Methods:
    _call: Executes a given prompt by invoking BRAD and optionally handles stop words and callback management.
    _llm_type: Returns the LLM type identifier for LangChain integration.
"""

from typing import Any, Dict, Iterator, List, Mapping, Optional
from langchain_core.language_models.llms import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun

class BradLLM(LLM):
    bot: Optional[Any] = None
    """
    A custom LLM class for integrating BRAD with LangChain. This class wraps BRAD, making
    it accessible through LangChain's LLM API, allowing BRAD to generate text completions 
    from prompts in LangChain-based applications.

    Attributes:
        bot (Optional[Any]): An instance of the BRAD chatbot that processes the prompts. 
            The bot should be set to an instance of BRAD externally before invoking the LLM.
    """
    
    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """
        Executes the LLM on a given prompt by passing it to the BRAD chatbot instance.
        
        :param prompt: The input prompt from which to generate text. This string is passed directly 
                       to BRAD for processing.
        :type prompt: str
        
        :param stop: A list of stop words to terminate text generation. If any stop word appears in 
                     the model's output, generation is halted. This feature is not implemented here, 
                     so passing this argument will trigger a warning message.
        :type stop: Optional[List[str]]
        
        :param run_manager: An optional callback manager for handling run-level events such as logging 
                            or tracing. Currently, this is not utilized in this implementation.
        :type run_manager: Optional[CallbackManagerForLLMRun]
        
        :param kwargs: Additional keyword arguments for model-specific behaviors or configurations.
        
        :raises NotImplementedError: If stop words are provided, as stopping criteria are not supported by BRAD.
        
        :return: The text generated by BRAD in response to the input prompt.
        :rtype: str
        """
        # Auth: Joshua Pickard
        #       jpic@umich.edu
        # Date: July 10, 2024

        if stop is not None:
            print("stop kwargs are not permitted.")
        return self.bot.invoke(prompt)
    
    @property
    def _llm_type(self) -> str:
        """
        Returns the type of the LLM, which is used to identify the model within LangChain.
        
        This property identifies the LLM as BRAD, allowing it to be registered and used 
        as a custom model in LangChain's workflows.
        
        :return: The string identifier for this LLM, which is "BRAD".
        :rtype: str
        """
        # Auth: Joshua Pickard
        #       jpic@umich.edu
        # Date: July 10, 2024

        return "BRAD"
