[AUR-411] Adopt to Example2 project (#28)

Add the chatbot from Example2. Create the UI for chat.

knowledgehub/chatbot/__init__.py

@@ -0,0 +1,4 @@
+from .base import BaseChatBot, ChatConversation
+from .simple_respondent import SimpleRespondentChatbot
+
+__all__ = ["BaseChatBot", "SimpleRespondentChatbot", "ChatConversation"]

knowledgehub/chatbot/base.py

@@ -0,0 +1,117 @@
+from abc import abstractmethod
+from typing import List, Optional
+
+from langchain.schema.messages import AIMessage, SystemMessage
+from theflow import Param, SessionCompose
+
+from ..base import BaseComponent
+from ..llms.base import LLMInterface
+from ..llms.chats.base import BaseMessage, HumanMessage
+
+
+class BaseChatBot(BaseComponent):
+    @abstractmethod
+    def run(self, messages: List[BaseMessage]) -> LLMInterface:
+        ...
+
+
+def session_chat_storage(obj):
+    """Store using the bot location rather than the session location"""
+    return obj._store_result
+
+
+class ChatConversation(SessionCompose):
+    """Base implementation of a chat bot component
+
+    A chatbot component should:
+        - handle internal state, including history messages
+        - return output for a given input
+    """
+
+    class Config:
+        store_result = session_chat_storage
+
+    system_message: Param[str] = Param(default="")
+    bot: BaseChatBot
+
+    def __init__(self, *args, **kwargs):
+        self._history: List[BaseMessage] = []
+        self._store_result = (
+            f"{self.__module__}.{self.__class__.__name__},uninitiated_bot"
+        )
+        super().__init__(*args, **kwargs)
+
+    def run(self, message: HumanMessage) -> Optional[BaseMessage]:
+        """Chat, given a message, return a response
+
+        Args:
+            message: The message to respond to
+
+        Returns:
+            The response to the message. If None, no response is sent.
+        """
+        user_message = (
+            HumanMessage(content=message) if isinstance(message, str) else message
+        )
+        self.history.append(user_message)
+
+        output = self.bot(self.history).text
+        output_message = None
+        if output is not None:
+            output_message = AIMessage(content=output)
+            self.history.append(output_message)
+
+        return output_message
+
+    def start_session(self):
+        self._store_result = self.bot.config.store_result
+        super().start_session()
+        if not self.history and self.system_message:
+            system_message = SystemMessage(content=self.system_message)
+            self.history.append(system_message)
+
+    def end_session(self):
+        super().end_session()
+        self._history = []
+
+    def check_end(
+        self,
+        history: Optional[List[BaseMessage]] = None,
+        user_message: Optional[HumanMessage] = None,
+        bot_message: Optional[AIMessage] = None,
+    ) -> bool:
+        """Check if a conversation should end"""
+        if user_message is not None and user_message.content == "":
+            return True
+
+        return False
+
+    def terminal_session(self):
+        """Create a terminal session"""
+        self.start_session()
+        print(">> Start chat:")
+
+        while True:
+            human = HumanMessage(content=input("Human: "))
+            if self.check_end(history=self.history, user_message=human):
+                break
+
+            output = self(human)
+            if output is None:
+                print("AI: <No response>")
+            else:
+                print("AI:", output.content)
+
+            if self.check_end(history=self.history, bot_message=output):
+                break
+
+        self.end_session()
+
+    @property
+    def history(self):
+        return self._history
+
+    @history.setter
+    def history(self, value):
+        self._history = value
+        self._variablex()

knowledgehub/chatbot/simple_respondent.py

@@ -0,0 +1,13 @@
+from theflow import Node
+
+from ..llms import ChatLLM
+from .base import BaseChatBot
+
+
+class SimpleRespondentChatbot(BaseChatBot):
+    """Simple text respondent chatbot that essentially wraps around a chat LLM"""
+
+    llm: Node[ChatLLM]
+
+    def _get_message(self) -> str:
+        return self.llm(self.history).text[0]