diff --git a/matrix_pygmalion_bot/bot/ai/langchain.py b/matrix_pygmalion_bot/bot/ai/langchain.py index a91afb7..9c0755b 100644 --- a/matrix_pygmalion_bot/bot/ai/langchain.py +++ b/matrix_pygmalion_bot/bot/ai/langchain.py @@ -1,7 +1,7 @@ import asyncio import os, time from .prompts import * -from .langchain_memory import CustomMemory # BotConversationSummaryBufferWindowMemory, TestMemory +from .langchain_memory import CustomMemory, ChangeNamesMemory # BotConversationSummaryBufferWindowMemory, TestMemory from ..utilities.messages import Message from langchain import PromptTemplate @@ -91,6 +91,8 @@ class AI(object): from ..wrappers.langchain_koboldcpp import KoboldCpp self.llm_chat = KoboldCpp(temperature=self.bot.temperature, endpoint_url="http://172.16.85.10:5001/api/latest/generate", stop=['<|endoftext|>']) self.llm_summary = KoboldCpp(temperature=0.2, endpoint_url="http://172.16.85.10:5002/api/latest/generate", stop=['<|endoftext|>'], max_tokens=512) + self.llm_chat_model = "pygmalion-7b" + self.llm_summary_model = "vicuna-13b" self.text_wrapper = text_wrapper self.image_wrapper = image_wrapper self.embeddings = SentenceTransformerEmbeddings() @@ -99,19 +101,21 @@ class AI(object): #self.memory = BotConversationSummerBufferWindowMemory(llm=self.llm_summary, max_token_limit=1200, min_token_limit=200) - def get_memory(self, room_id, human_prefix="Human"): + def get_memory(self, room_id, human_prefix=None): if not room_id in self.rooms: self.rooms[room_id] = {} if "moving_summary" in self.bot.rooms[room_id]: moving_summary = self.bot.rooms[room_id]['moving_summary'] else: moving_summary = "No previous events." + if not human_prefix: + human_prefix = "Human" memory = CustomMemory(memory_key="chat_history", input_key="input", human_prefix=human_prefix, ai_prefix=self.bot.name, llm=self.llm_summary, summary_prompt=prompt_progressive_summary, moving_summary_buffer=moving_summary, max_len=1200, min_len=200) self.rooms[room_id]["memory"] = memory #memory.chat_memory.add_ai_message(self.bot.greeting) else: memory = self.rooms[room_id]["memory"] - if human_prefix != memory.human_prefix: + if human_prefix: memory.human_prefix = human_prefix return memory @@ -185,12 +189,22 @@ class AI(object): chat_ai_name = self.bot.name chat_human_name = message.additional_kwargs['user_name'] room_id = message.additional_kwargs['room_id'] - if False: # model is vicuna + + if self.llm_chat_model.startswith('vicuna'): + prompt_chat = prompt_vicuna chat_ai_name = "### Assistant" chat_human_name = "### Human" + elif self.llm_chat_model.startswith('pygmalion'): + prompt_chat = prompt_pygmalion + chat_human_name = "You" + elif self.llm_chat_model.startswith('koboldai'): + prompt_chat = prompt_koboldai + else: + prompt_chat = prompt_alpaca conversation_memory = self.get_memory(room_id, chat_human_name) readonlymemory = ReadOnlySharedMemory(memory=conversation_memory) + custom_memory = ChangeNamesMemory(memory=conversation_memory, replace_ai_chat_names={self.bot.name: chat_ai_name}, replace_human_chat_names={message.additional_kwargs['user_name']: chat_human_name}) #summary_memory = ConversationSummaryMemory(llm=self.llm_summary, memory_key="summary", input_key="input") #combined_memory = CombinedMemory(memories=[conversation_memory, summary_memory]) @@ -200,20 +214,23 @@ class AI(object): #when = humanize.naturaltime(t) #print(when) - # ToDo: either use prompt.format() to fill out the pygmalion prompt and use # the resulting template text to feed it into the instruct prompt's instruction # or do this with the prompt.partial() - prompt = prompt_vicuna.partial( + prompt = prompt_chat.partial( ai_name=self.bot.name, persona=self.bot.persona, scenario=self.bot.scenario, - summary=conversation_memory.moving_summary_buffer, human_name=chat_human_name, - #example_dialogue=replace_all(self.bot.example_dialogue, {"{{user}}": chat_human_name, "{{char}}": chat_ai_name}) ai_name_chat=chat_ai_name, ) + if "summary" in prompt_chat.input_variables: + prompt = prompt.partial(summary=conversation_memory.moving_summary_buffer) + if "example_dialogue" in prompt_chat.input_variables: + prompt = prompt.partial( + example_dialogue=self.bot.example_dialogue.replace("{{user}}", chat_human_name) + ) tmp_prompt_text = prompt.format(chat_history=conversation_memory.buffer, input=message.content) prompt_len = self.llm_chat.get_num_tokens(tmp_prompt_text) @@ -230,7 +247,7 @@ class AI(object): llm=self.llm_chat, prompt=prompt, verbose=True, - memory=readonlymemory, + memory=custom_memory, #stop=['<|endoftext|>', '\nYou:', f"\n{chat_human_name}:"], ) @@ -292,7 +309,7 @@ class AI(object): else: input_text = conversation_memory.moving_summary_buffer - return await diary_chain.apredict(text=input_text) + return await diary_chain.apredict(text=input_text, ai_name=self.bot.name) async def agent(self): @@ -358,6 +375,7 @@ class AI(object): async def sleep(self): + logger.info(f"{self.bot.name} sleeping now... running background tasks...") # Write Date into chat history for room_id in self.rooms.keys(): #fake_message = Message(datetime.now().timestamp(), self.bot.name, "", event_id=None, user_id=None, room_name=None, room_id=room_id) @@ -389,7 +407,9 @@ class AI(object): # Update stats # Let background tasks run conversation_memory.chat_memory_day.clear() + await conversation_memory.prune_memory(conversation_memory.min_len) await self.bot.write_conf2(self.bot.rooms) + logger.info(f"{self.bot.name} done sleeping and ready for the next day...") async def prime_llm(self, text): @@ -397,6 +417,7 @@ class AI(object): def replace_all(text, dic): + #example_dialogue=replace_all(self.bot.example_dialogue, {"{{user}}": chat_human_name, "{{char}}": chat_ai_name}) for i, j in dic.items(): text = text.replace(i, j) return text diff --git a/matrix_pygmalion_bot/bot/ai/langchain_memory.py b/matrix_pygmalion_bot/bot/ai/langchain_memory.py index ed3abce..32103b3 100644 --- a/matrix_pygmalion_bot/bot/ai/langchain_memory.py +++ b/matrix_pygmalion_bot/bot/ai/langchain_memory.py @@ -97,7 +97,7 @@ class CustomMemory(BaseMemory): curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer) if curr_buffer_length > max_len: pruned_memory = [] - while curr_buffer_length > self.min_len and len(buffer) > 0: + while curr_buffer_length > min(self.min_len, self.max_len) and len(buffer) > 0: pruned_memory.append(buffer.pop(0)) curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer) self.moving_summary_buffer = await self.apredict_new_summary(pruned_memory, self.moving_summary_buffer) @@ -130,6 +130,8 @@ class CustomMemory(BaseMemory): role = m.role else: raise ValueError(f"Got unsupported message type: {m}") + if "user_name" in m.additional_kwargs: + role = m.additional_kwargs["user_name"] string_messages.append(f"{role}: {m.content}") return "\n".join(string_messages) @@ -181,6 +183,66 @@ class CustomMemory(BaseMemory): return output.strip() +class ChangeNamesMemory(BaseMemory): + """A memory wrapper that changes names.""" + + memory: BaseMemory + replace_ai_chat_names: Dict[str, str] + replace_human_chat_names: Dict[str, str] + + def get_buffer_string(self, messages: List[BaseMessage]) -> str: + """Get buffer string of messages.""" + string_messages = [] + for m in messages: + if isinstance(m, HumanMessage): + role = self.memory.human_prefix + elif isinstance(m, AIMessage): + role = self.memory.ai_prefix + elif isinstance(m, SystemMessage): + role = "System" + elif isinstance(m, ChatMessage): + role = m.role + else: + raise ValueError(f"Got unsupported message type: {m}") + if "user_name" in m.additional_kwargs: + role = m.additional_kwargs["user_name"] + if isinstance(m, HumanMessage): + for i, j in self.replace_human_chat_names.items(): + role = role.replace(i, j) + elif isinstance(m, AIMessage): + for i, j in self.replace_ai_chat_names.items(): + role = role.replace(i, j) + string_messages.append(f"{role}: {m.content}") + return "\n".join(string_messages) + + @property + def buffer(self) -> Any: + """String buffer of memory.""" + if self.memory.return_messages: + return self.memory.chat_memory.messages + else: + return self.get_buffer_string( + self.memory.chat_memory.messages, + ) + + @property + def memory_variables(self) -> List[str]: + """Return memory variables.""" + return self.memory.memory_variables + + def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]: + """Load memory variables from memory.""" + return {self.memory.memory_key: self.buffer} + + def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None: + """Nothing should be saved or changed""" + pass + + def clear(self) -> None: + """Nothing to clear, got a memory like a vault.""" + pass + + class ChatMessageHistoryMessage(BaseModel): #messages: List[Message] = [] messages = [] diff --git a/matrix_pygmalion_bot/bot/ai/prompts.py b/matrix_pygmalion_bot/bot/ai/prompts.py index 87a4b4f..99c2c43 100644 --- a/matrix_pygmalion_bot/bot/ai/prompts.py +++ b/matrix_pygmalion_bot/bot/ai/prompts.py @@ -4,10 +4,11 @@ from langchain import PromptTemplate prompt_pygmalion = PromptTemplate.from_template( """{ai_name}'s Persona: {persona} Scenario: {scenario} +Summary of previous events: {summary} {chat_history} -{human_name}: {human_input} +{human_name}: {input} {ai_name_chat}:""" ) @@ -16,7 +17,7 @@ prompt_koboldai = PromptTemplate.from_template( [Start Scene: {scenario}] {chat_history} -{human_name}: {human_input} +{human_name}: {input} {ai_name_chat}:""" ) @@ -31,10 +32,10 @@ Scenario: {scenario} ### Response: {chat_history} -{human_name}: {human_input} +{human_name}: {input} {ai_name_chat}:""" prompt_alpaca = PromptTemplate( - input_variables=["ai_name", "persona", "scenario", "chat_history", "human_name", "ai_name_chat", "human_input"], + input_variables=["ai_name", "persona", "scenario", "chat_history", "human_name", "ai_name_chat", "input"], template=template_alpaca, ) @@ -134,7 +135,7 @@ prompt_outline = PromptTemplate.from_template( """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. ### Instruction: -Provide an outline of the character's day in keywords. +Provide an outline of {ai_name}'s day in keywords. ### Input: {text} diff --git a/matrix_pygmalion_bot/bot/utilities/messages.py b/matrix_pygmalion_bot/bot/utilities/messages.py index 6166b1b..801b1ce 100644 --- a/matrix_pygmalion_bot/bot/utilities/messages.py +++ b/matrix_pygmalion_bot/bot/utilities/messages.py @@ -20,7 +20,6 @@ class Message(object): if self.role == "human": return HumanMessage( content=self.message, - role=self.user_name, # "chat" additional_kwargs={ "timestamp": self.timestamp, "user_name": self.user_name, @@ -33,7 +32,6 @@ class Message(object): elif self.role == "ai": return AIMessage( content=self.message, - role=self.user_name, # "chat" additional_kwargs={ "timestamp": self.timestamp, "user_name": self.user_name, @@ -46,7 +44,6 @@ class Message(object): elif self.role == "system": return SystemMessage( content=self.message, - role=self.user_name, # "chat" additional_kwargs={ "timestamp": self.timestamp, "user_name": self.user_name,