∞ Chat Engine

Documentation

Class name: LLMChatBot
Category: SALT/Language Toolkit/Querying
Output node: False

The LLMChatBot node is designed to facilitate interactive chat sessions using large language models (LLMs). It leverages a combination of user inputs, context, and predefined settings to generate conversational responses, aiming to simulate a natural and engaging dialogue experience.

Input types

Required

llm_model
- Specifies the large language model to be used for the chat session, including any relevant configurations and settings.
- Comfy dtype: LLM_MODEL
- Python dtype: Dict[str, Any]
llm_context
- Provides additional context or settings specific to the chat session, enhancing the model's understanding and response accuracy.
- Comfy dtype: LLM_CONTEXT
- Python dtype: Any
prompt
- The user's input message or question to which the chatbot will respond, serving as the trigger for generating a reply.
- Comfy dtype: STRING
- Python dtype: str

Optional

reset_engine
- A flag to reset the chat engine, useful for starting a new conversation thread or clearing history.
- Comfy dtype: BOOLEAN
- Python dtype: bool
user_nickname
- A customizable nickname for the user, used to personalize the chat session.
- Comfy dtype: STRING
- Python dtype: str
system_nickname
- A customizable nickname for the chatbot, enhancing the conversational experience by adding a layer of personalization.
- Comfy dtype: STRING
- Python dtype: str
char_per_token
- Defines the average number of characters per token, used for tokenization purposes in the chat session.
- Comfy dtype: INT
- Python dtype: int

Output types

chat_history
- Comfy dtype: STRING
- The accumulated conversation history between the user and the chatbot, including all exchanged messages.
- Python dtype: List[Dict[str, Any]]
response
- Comfy dtype: STRING
- The chatbot's direct response to the user's latest prompt.
- Python dtype: str
chat_token_count
- Comfy dtype: INT
- The total number of tokens used in the chat session, reflecting the conversation's complexity and length.
- Python dtype: int

Usage tips

Infra type: CPU
Common nodes: unknown

Source code

class LLMChatBot:
    def __init__(self):
        self.chat_history = []
        self.history = []
        self.token_map = {} 

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "llm_model": ("LLM_MODEL", ),
                "llm_context": ("LLM_CONTEXT", ),
                "prompt": ("STRING", {"multiline": True, "dynamicPrompt": False}),
            },
            "optional": {
                "reset_engine": ("BOOLEAN", {"default": False}),
                "user_nickname": ("STRING", {"default": "User"}),
                "system_nickname": ("STRING", {"default": "Assistant"}),
                "char_per_token": ("INT", {"min": 1, "default": 4})
            }
        }

    RETURN_TYPES = ("STRING", "STRING", "INT")
    RETURN_NAMES = ("chat_history", "response", "chat_token_count")

    FUNCTION = "chat"
    CATEGORY = f"{MENU_NAME}/{SUB_MENU_NAME}/Querying"

    def chat(self, llm_model: Dict[str, Any], llm_context:Any, prompt: str, reset_engine:bool = False, user_nickname:str = "User", system_nickname:str = "Assistant", char_per_token:int = 4) -> str:

        if reset_engine:
            self.chat_history.clear()
            self.history.clear()
            self.token_map.clear()

        max_tokens = llm_model.get("max_tokens", 4096)
        using_mock_tokenizer = False
        try:
            tokenizer = tiktoken.encoding_for_model(llm_model.get('llm_name', 'gpt-3-turbo'))
        except Exception:
            using_mock_tokenizer = True
            tokenizer = MockTokenizer(max_tokens, char_per_token=char_per_token)

        if not self.chat_history:
            system_prompt = getattr(llm_model['llm'], "system_prompt", None)
            if system_prompt not in (None, ""):
                initial_msg = ChatMessage(role=MessageRole.SYSTEM, content=system_prompt)
                self.chat_history.append(initial_msg)
                self.token_map[0] = tokenizer.encode(system_prompt)

        # Tokenize and count initial tokens
        cumulative_token_count = 0
        for index, message in enumerate(self.chat_history):
            if index not in self.token_map:
                self.token_map[index] = tokenizer.encode(message.content)
            #if not using_mock_tokenizer:
            cumulative_token_count += len(self.token_map[index])
            #else:
            #    cumulative_token_count += tokenizer.count(self.token_map[index])

        # Prune messages from the history if over max_tokens
        index = 0
        while cumulative_token_count > max_tokens and index < len(self.chat_history):
            tokens = self.token_map[index]
            token_count = len(tokens) #if not using_mock_tokenizer else tokenizer.count(tokens)
            if token_count > 1:
                tokens.pop(0)
                self.chat_history[index].content = tokenizer.decode(tokens)
                cumulative_token_count -= 1
            else:
                cumulative_token_count -= token_count
                self.chat_history.pop(index)
                self.token_map.pop(index)
                for old_index in list(self.token_map.keys()):
                    if old_index > index:
                        self.token_map[old_index - 1] = self.token_map.pop(old_index)
                continue
            index += 1
                
        history_string = ""
        reply_string = ""
        documents = []

        # Build prior history string
        for history in self.history:
            user, assistant, timestamp = history
            history_string += f"""[{user_nickname}]: {history[user]}

[{system_nickname}]: {history[assistant]}

"""
        # Spoof documents -- Why can't we just talk to a modeL?
        documents = [Document(text="null", extra_info={})]

        index = VectorStoreIndex.from_documents(
            documents, 
            service_context=llm_context,
            transformations=[SentenceSplitter(chunk_size=1024, chunk_overlap=20)]
        )
        chat_engine = index.as_chat_engine(chat_mode="best")

        response = chat_engine.chat(prompt, chat_history=self.chat_history)

        response_dict = {
            user_nickname: prompt, 
            system_nickname: response.response,
            "timestamp": str(time.time())
        }

        user_cm = ChatMessage(role=MessageRole.USER, content=prompt)
        system_cm = ChatMessage(role=MessageRole.SYSTEM, content=response.response)
        self.chat_history.append(user_cm)
        self.chat_history.append(system_cm)

        self.history.append(response_dict)

        reply_string = response.response

        history_string += f"""[{user_nickname}]: {prompt}

[{system_nickname}]: {response.response}"""
            
        pprint(self.chat_history, indent=4)

        return (history_string, reply_string, cumulative_token_count)

7.1 KiB Raw Permalink Blame History