| from data import Data | |
| ''' | |
| The BotWrapper class makes it so that different types of bots can be used in the same way. | |
| This is used in the Bots class to create a list of all bots and pass them to the frontend. | |
| ''' | |
| class BotWrapper: | |
| def __init__(self, bot): | |
| self.bot = bot | |
| def chat(self, *args, **kwargs): | |
| methods = ['chat', 'query'] | |
| for method in methods: | |
| if hasattr(self.bot, method): | |
| print(f"Calling {method} method") | |
| method_to_call = getattr(self.bot, method) | |
| return method_to_call(*args, **kwargs).response() | |
| raise AttributeError(f"'{self.bot.__class__.__name__}' object has none of the required methods: '{methods}'") | |
| def stream_chat(self, *args, **kwargs): | |
| methods = ['stream_chat', 'query'] | |
| for method in methods: | |
| if hasattr(self.bot, method): | |
| print(f"Calling {method} method") | |
| method_to_call = getattr(self.bot, method) | |
| return method_to_call(*args, **kwargs).response_gen | |
| raise AttributeError(f"'{self.bot.__class__.__name__}' object has none of the required methods: '{methods}'") | |
| ''' | |
| The Bots class creates the bots and passes them to the frontend. | |
| ''' | |
| class Bots: | |
| def __init__(self): | |
| self.data = Data() | |
| self.data.load_data() | |
| self.query_engine = None | |
| self.chat_agent = None | |
| self.all_bots = None | |
| self.create_bots() | |
| def create_query_engine_bot(self): | |
| if self.query_engine is None: | |
| self.query_engine = BotWrapper(self.data.index.as_query_engine()) | |
| return self.query_engine | |
| def create_chat_agent(self): | |
| if self.chat_agent is None: | |
| from llama_index.core.memory import ChatMemoryBuffer | |
| memory = ChatMemoryBuffer.from_defaults(token_limit=1500) | |
| self.chat_agent = BotWrapper(self.data.index.as_chat_engine( | |
| chat_mode="context", | |
| memory=memory, | |
| context_prompt=( | |
| "You are a chatbot, able to have normal interactions, as well as talk" | |
| " about the questions and answers you know about." | |
| "Here are the relevant documents for the context:\n" | |
| "{context_str}" | |
| "\nInstruction: Use the previous chat history, or the context above, to interact and help the user." | |
| ) | |
| )) | |
| return self.chat_agent | |
| def create_bots(self): | |
| self.create_query_engine_bot() | |
| self.create_chat_agent() | |
| self.all_bots = [self.query_engine, self.chat_agent] | |
| return self.all_bots | |
| def get_bots(self): | |
| return self.all_bots | |