File size: 2,752 Bytes
60d9d3a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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
|