From ee1ffbac6bbb526bea1598e670aa4688f183fa56 Mon Sep 17 00:00:00 2001 From: Hendrik Langer Date: Fri, 5 May 2023 16:01:12 +0200 Subject: [PATCH] more work on langchain memory --- matrix_pygmalion_bot/bot/ai/langchain.py | 81 ++++----- .../bot/ai/langchain_memory.py | 167 +++++++++++++++++- matrix_pygmalion_bot/bot/ai/prompts.py | 1 + matrix_pygmalion_bot/bot/core.py | 9 +- .../bot/utilities/messages.py | 67 +++++-- .../bot/wrappers/langchain_koboldcpp.py | 7 +- 6 files changed, 261 insertions(+), 71 deletions(-) diff --git a/matrix_pygmalion_bot/bot/ai/langchain.py b/matrix_pygmalion_bot/bot/ai/langchain.py index fe481d8..55166a9 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 BotConversationSummaryBufferWindowMemory, TestMemory +from .langchain_memory import CustomMemory # BotConversationSummaryBufferWindowMemory, TestMemory from ..utilities.messages import Message from langchain import PromptTemplate @@ -90,7 +90,7 @@ 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:5001/api/latest/generate", stop=['<|endoftext|>'], max_tokens=512) + self.llm_summary = KoboldCpp(temperature=0.2, endpoint_url="http://172.16.85.10:5002/api/latest/generate", stop=['<|endoftext|>'], max_tokens=512) self.text_wrapper = text_wrapper self.image_wrapper = image_wrapper self.embeddings = SentenceTransformerEmbeddings() @@ -102,7 +102,7 @@ class AI(object): def get_memory(self, room_id, human_prefix="Human"): if not room_id in self.rooms: self.rooms[room_id] = {} - memory = ConversationBufferMemory(memory_key="chat_history", input_key="input", human_prefix=human_prefix, ai_prefix=self.bot.name) + 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, max_len=1200, min_len=200) self.rooms[room_id]["memory"] = memory self.rooms[room_id]["summary"] = "No previous events." memory.chat_memory.add_ai_message(self.bot.greeting) @@ -181,25 +181,15 @@ class AI(object): async def generate_roleplay(self, message, reply_fn, typing_fn): - langchain_human_message = HumanMessage( - content=message.message, - additional_kwargs={ - "timestamp": message.timestamp, - "user_name": message.user_name, - "event_id": message.event_id, - "user_id": message.user_id, - "room_name": message.room_name, - "room_id": message.room_id, - } - ) chat_ai_name = self.bot.name - chat_human_name = message.user_name + chat_human_name = message.additional_kwargs['user_name'] + room_id = message.additional_kwargs['room_id'] if False: # model is vicuna chat_ai_name = "### Assistant" chat_human_name = "### Human" - conversation_memory = self.get_memory(message.room_id, message.user_name) + conversation_memory = self.get_memory(room_id, chat_human_name) conversation_memory.human_prefix = chat_human_name readonlymemory = ReadOnlySharedMemory(memory=conversation_memory) summary_memory = ConversationSummaryMemory(llm=self.llm_summary, memory_key="summary", input_key="input") @@ -211,11 +201,11 @@ class AI(object): async def make_progressive_summary(previous_summary, chat_history_text_string): await asyncio.sleep(0) # yield for matrix-nio - #self.rooms[message.room_id]["summary"] = summary_memory.predict_new_summary(conversation_memory.chat_memory.messages, previous_summary).strip() + #self.rooms[room_id]["summary"] = summary_memory.predict_new_summary(conversation_memory.chat_memory.messages, previous_summary).strip() summary_chain = LLMChain(llm=self.llm_summary, prompt=prompt_progressive_summary, verbose=True) - self.rooms[message.room_id]["summary"] = await summary_chain.apredict(summary=previous_summary, chat_history=chat_history_text_string) + self.rooms[room_id]["summary"] = await summary_chain.apredict(summary=previous_summary, chat_history=chat_history_text_string) # ToDo: maybe add an add_task_done callback and don't access the variable directly from here? - logger.info(f"New summary is: \"{self.rooms[message.room_id]['summary']}\"") + logger.info(f"New summary is: \"{self.rooms[room_id]['summary']}\"") conversation_memory.chat_memory.messages = conversation_memory.chat_memory.messages[-k * 2 :] conversation_memory.load_memory_variables({}) #summary = summarize(conversation_memory.buffer) @@ -224,11 +214,11 @@ class AI(object): logger.info("memory progressive summary scheduled...") - await self.bot.schedule(self.bot.queue, make_progressive_summary, self.rooms[message.room_id]["summary"], conversation_memory.buffer) #.add_done_callback( + await self.bot.schedule(self.bot.queue, make_progressive_summary, self.rooms[room_id]["summary"], conversation_memory.buffer) #.add_done_callback( - #t = datetime.fromtimestamp(message.timestamp) + #t = datetime.fromtimestamp(message.additional_kwargs['timestamp']) #when = humanize.naturaltime(t) #print(when) @@ -241,8 +231,8 @@ class AI(object): ai_name=self.bot.name, persona=self.bot.persona, scenario=self.bot.scenario, - summary=self.rooms[message.room_id]["summary"], - human_name=message.user_name, + summary=self.rooms[room_id]["summary"], + 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, ) @@ -252,48 +242,44 @@ class AI(object): prompt=prompt, verbose=True, memory=readonlymemory, - #stop=['<|endoftext|>', '\nYou:', f"\n{message.user_name}:"], + #stop=['<|endoftext|>', '\nYou:', f"\n{chat_human_name}:"], ) -# output = llm_chain(inputs={"ai_name": self.bot.name, "persona": self.bot.persona, "scenario": self.bot.scenario, "human_name": message.user_name, "ai_name_chat": self.bot.name, "chat_history": "", "input": message.message})['results'][0]['text'] +# output = llm_chain(inputs={"ai_name": self.bot.name, "persona": self.bot.persona, "scenario": self.bot.scenario, "human_name": chat_human_name, "ai_name_chat": self.bot.name, "chat_history": "", "input": message.content})['results'][0]['text'] #roleplay_chain = RoleplayChain(llm_chain=chain, character_name=self.bot.name, persona=self.bot.persona, scenario=self.bot.scenario, ai_name_chat=chat_ai_name, human_name_chat=chat_human_name) stop = ['<|endoftext|>', f"\n{chat_human_name}"] - #print(f"Message is: \"{message.message}\"") + #print(f"Message is: \"{message.content}\"") await asyncio.sleep(0) - output = await chain.arun({"input":message.message, "stop": stop}) - output = output.replace("", self.bot.name).replace("", message.user_name) + output = await chain.arun({"input":message.content, "stop": stop}) + output = output.replace("", self.bot.name).replace("", chat_human_name) output = output.replace("### Assistant", self.bot.name) output = output.replace(f"\n{self.bot.name}: ", " ") output = output.strip() + + if "*activates the neural uplink*" in output.casefold(): + pass # call agent + + own_message_resp = await reply_fn(output) + langchain_ai_message = AIMessage( content=output, additional_kwargs={ "timestamp": datetime.now().timestamp(), "user_name": self.bot.name, - "event_id": None, + "event_id": own_message_resp.event_id, "user_id": None, - "room_name": message.room_name, - "room_id": message.room_id, + "room_name": message.additional_kwargs['room_name'], + "room_id": own_message_resp.room_id, } ) - if "*activates the neural uplink*" in output.casefold(): - pass # call agent - - #conversation_memory.chat_memory.messages.append(ChatMessage(content=message, role=message.user_name)) - conversation_memory.chat_memory.add_user_message(message.message) - conversation_memory.chat_memory.add_ai_message(output) + conversation_memory.save_context({"input": message.content}, {"ouput": output}) conversation_memory.load_memory_variables({}) - if not "messages_today" in self.rooms[message.room_id]: - self.rooms[message.room_id]["messages_today"] = [] - self.rooms[message.room_id]["messages_today"].append(langchain_human_message) - self.rooms[message.room_id]["messages_today"].append(langchain_ai_message) - - return output.strip() + return output async def summarize(self, text): @@ -306,10 +292,11 @@ class AI(object): async def diary(self, room_id): await asyncio.sleep(0) # yield for matrix-nio diary_chain = LLMChain(llm=self.llm_summary, prompt=prompt_outline, verbose=True) + conversation_memory = self.get_memory(room_id) #self.rooms[message.room_id]["summary"] string_messages = [] - for m in self.rooms[room_id]["messages_today"]: - string_messages.append(f"{message.user_name}: {message.message}") + for m in conversation_memory.chat_memory_day.messages: + string_messages.append(f"{message.role}: {message.content}") return await diary_chain.apredict(text="\n".join(string_messages)) @@ -397,12 +384,12 @@ class AI(object): # Summarize the last day and save a diary entry yesterday = ( datetime.now() - timedelta(days=1) ).strftime('%Y-%m-%d') for room_id in self.rooms.keys(): - if "messages_today" in self.rooms[room_id]: + if len(conversation_memory.chat_memory_day.messages) > 0: self.bot.rooms[room_id]["diary"][yesterday] = await self.diary(room_id) # Calculate new goals for the character # Update stats # Let background tasks run - self.rooms[room_id]["messages_today"] = [] + conversation_memory.chat_memory_day.clear() await self.bot.write_conf2(self.bot.rooms) diff --git a/matrix_pygmalion_bot/bot/ai/langchain_memory.py b/matrix_pygmalion_bot/bot/ai/langchain_memory.py index 719c89f..be5a4e8 100644 --- a/matrix_pygmalion_bot/bot/ai/langchain_memory.py +++ b/matrix_pygmalion_bot/bot/ai/langchain_memory.py @@ -1,15 +1,170 @@ -from typing import Any, Dict, List +import asyncio +from typing import Any, Dict, List, Tuple, Optional +from pydantic import BaseModel, Extra, Field, root_validator from langchain.chains.llm import LLMChain -from langchain.memory.chat_memory import BaseChatMemory +from langchain.memory.chat_memory import BaseChatMemory, BaseMemory from langchain.memory.prompt import SUMMARY_PROMPT from langchain.prompts.base import BasePromptTemplate -from langchain.schema import BaseLanguageModel, BaseMessage, get_buffer_string +from langchain.schema import BaseLanguageModel, BaseMessage, BaseChatMessageHistory, BaseMemory, get_buffer_string +from langchain.schema import AIMessage, HumanMessage, SystemMessage, ChatMessage +from langchain.memory.chat_message_histories.in_memory import ChatMessageHistory from ..utilities.messages import Message -class ChatMessageHistory(BaseModel): - messages: List[Message] = [] + +class ChatMessageHistoryCustom(BaseChatMessageHistory, BaseModel): + messages: List[BaseMessage] = [] + + def add_user_message(self, message: str) -> None: + self.messages.append(HumanMessage(content=message)) + + def add_ai_message(self, message: str) -> None: + self.messages.append(AIMessage(content=message)) + + def add_system_message(self, message: str) -> None: + self.messages.append(SystemMessage(content=message)) + + def add_chat_message(self, message: str) -> None: + self.messages.append(ChatMessage(content=message)) + + def clear(self) -> None: + self.messages = [] + + +class CustomMemory(BaseMemory): + """Buffer for storing conversation memory.""" + + human_prefix: str = "Human" + ai_prefix: str = "AI" + memory_key: str = "history" #: :meta private: + + chat_memory: BaseChatMessageHistory = Field(default_factory=ChatMessageHistoryCustom) + chat_memory_day: BaseChatMessageHistory = Field(default_factory=ChatMessageHistoryCustom) + output_key: Optional[str] = None + input_key: Optional[str] = None + return_messages: bool = False + + max_len: int = 1200 + min_len: int = 200 + #length_function: Callable[[str], int] = len, + #length_function: Callable[[str], int] = self.llm.get_num_tokens_from_messages, + + moving_summary_buffer: str = "" + + llm: BaseLanguageModel + summary_prompt: BasePromptTemplate = SUMMARY_PROMPT + #summary_message_cls: Type[BaseMessage] = SystemMessage + + def _get_input_output( + self, inputs: Dict[str, Any], outputs: Dict[str, str] + ) -> Tuple[str, str]: + if self.input_key is None: + prompt_input_key = get_prompt_input_key(inputs, self.memory_variables) + else: + prompt_input_key = self.input_key + if self.output_key is None: + if len(outputs) != 1: + raise ValueError(f"One output key expected, got {outputs.keys()}") + output_key = list(outputs.keys())[0] + else: + output_key = self.output_key + return inputs[prompt_input_key], outputs[output_key] + + def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None: + """Save context from this conversation to buffer.""" + input_str, output_str = self._get_input_output(inputs, outputs) + self.chat_memory.add_user_message(input_str) + self.chat_memory.add_ai_message(output_str) + self.chat_memory_day.add_user_message(input_str) + self.chat_memory_day.add_ai_message(output_str) + # Prune buffer if it exceeds max token limit + buffer = self.chat_memory.messages + curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer) + if curr_buffer_length > self.max_len: + pruned_memory = [] + while curr_buffer_length > self.min_len: + pruned_memory.append(buffer.pop(0)) + curr_buffer_length = self.llm.get_num_tokens_from_messages(buffer) + loop = asyncio.get_event_loop() + self.moving_summary_buffer = loop.run_until_complete( + self.apredict_new_summary(pruned_memory, self.moving_summary_buffer) + ) + + + def clear(self) -> None: + """Clear memory contents.""" + self.chat_memory.clear() + self.chat_memory_day.clear() + self.moving_summary_buffer = "" + + def get_buffer_string(self, messages: List[BaseMessage], human_prefix: str = "Human", ai_prefix: str = "AI") -> str: + """Get buffer string of messages.""" + string_messages = [] + for m in messages: + if isinstance(m, HumanMessage): + role = human_prefix + elif isinstance(m, AIMessage): + role = ai_prefix + elif isinstance(m, SystemMessage): + role = "System" + elif isinstance(m, ChatMessage): + role = m.role + else: + raise ValueError(f"Got unsupported message type: {m}") + string_messages.append(f"{role}: {m.content}") + return "\n".join(string_messages) + + @property + def buffer(self) -> Any: + """String buffer of memory.""" + if self.return_messages: + return self.chat_memory.messages + else: + return self.get_buffer_string( + self.chat_memory.messages, + human_prefix=self.human_prefix, + ai_prefix=self.ai_prefix, + ) + + @property + def buffer_day(self) -> Any: + """String buffer of memory.""" + if self.return_messages: + return self.chat_memory_day.messages + else: + return self.get_buffer_string( + self.chat_memory_day.messages, + human_prefix=self.human_prefix, + ai_prefix=self.ai_prefix, + ) + + @property + def memory_variables(self) -> List[str]: + """Will always return list of memory variables. + :meta private: + """ + return [self.memory_key] + + def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + """Return history buffer.""" + return {self.memory_key: self.buffer} + + async def apredict_new_summary(self, messages: List[BaseMessage], existing_summary: str) -> str: + new_lines = self.get_buffer_string( + messages, + human_prefix=self.human_prefix, + ai_prefix=self.ai_prefix, + ) + + chain = LLMChain(llm=self.llm, prompt=self.summary_prompt) + await asyncio.sleep(0) + return await chain.apredict(summary=existing_summary, new_lines=new_lines) + + +class ChatMessageHistoryMessage(BaseModel): + #messages: List[Message] = [] + messages = [] def add_user_message(self, message: Message) -> None: self.messages.append(message) @@ -32,7 +187,7 @@ class TestMemory(BaseMemory): human_prefix: str = "Human" ai_prefix: str = "AI" - chat_memory: ChatMessageHistory = Field(default_factory=ChatMessageHistory) + chat_memory: ChatMessageHistory = Field(default_factory=ChatMessageHistoryMessage) # buffer: str = "" output_key: Optional[str] = None input_key: Optional[str] = None diff --git a/matrix_pygmalion_bot/bot/ai/prompts.py b/matrix_pygmalion_bot/bot/ai/prompts.py index 5f4115c..87a4b4f 100644 --- a/matrix_pygmalion_bot/bot/ai/prompts.py +++ b/matrix_pygmalion_bot/bot/ai/prompts.py @@ -128,6 +128,7 @@ New summary: """ ) #Progressively summarize the lines of conversation provided, adding onto the previous summary returning a new summary. +#only include relevant facts for {{char}}'s long-term memory / future 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. diff --git a/matrix_pygmalion_bot/bot/core.py b/matrix_pygmalion_bot/bot/core.py index 5205cf2..b084682 100644 --- a/matrix_pygmalion_bot/bot/core.py +++ b/matrix_pygmalion_bot/bot/core.py @@ -154,9 +154,12 @@ class ChatBot(object): if message.is_from(self.connection.user_id): """Skip messages from ouselves""" + message.role = "ai" self.chatlog.save(message) await self.connection.room_read_markers(room.room_id, event.event_id, event.event_id) return + else: + message.role = "human" # if event.decrypted: # encrypted_symbol = "🛡 " @@ -191,6 +194,8 @@ class ChatBot(object): await self.schedule(self.queue, self.process_command, message, reply_fn, typing_fn) # elif re.search("^(?=.*\bsend\b)(?=.*\bpicture\b).*$", event.body, flags=re.IGNORECASE): # # send, mail, drop, snap picture, photo, image, portrait + elif message.is_error(): + return else: await self.schedule(self.queue, self.process_message, message, reply_fn, typing_fn) self.rooms[room.room_id]['num_messages'] += 1 @@ -257,10 +262,10 @@ class ChatBot(object): async def process_message(self, message, reply_fn, typing_fn): - output = await self.ai.generate_roleplay(message, reply_fn, typing_fn) + output = await self.ai.generate_roleplay(message.to_langchain(), reply_fn, typing_fn) #output = await self.ai.generate(message, reply_fn, typing_fn) # typing false - await reply_fn(output) + #await reply_fn(output) async def event_loop(self): diff --git a/matrix_pygmalion_bot/bot/utilities/messages.py b/matrix_pygmalion_bot/bot/utilities/messages.py index 70721a1..6ef3f28 100644 --- a/matrix_pygmalion_bot/bot/utilities/messages.py +++ b/matrix_pygmalion_bot/bot/utilities/messages.py @@ -2,7 +2,7 @@ from langchain.schema import AIMessage, HumanMessage, SystemMessage, ChatMessage class Message(object): - def __init__(self, timestamp, user_name, message, event_id=None, user_id=None, room_name=None, room_id=None): + def __init__(self, timestamp, user_name, message, event_id=None, user_id=None, room_name=None, room_id=None, role=None): self.timestamp = timestamp self.user_name = user_name self.message = message @@ -10,24 +10,65 @@ class Message(object): self.user_id = user_id self.room_name = room_name self.room_id = room_id + self.role = role @classmethod def from_matrix(cls, room, event): return cls(event.server_timestamp/1000, room.user_name(event.sender), event.body, event.event_id, event.sender, room.display_name, room.room_id) def to_langchain(self): - return ChatMessage( - content=self.message, - role=self.user_name, # "chat" - additional_kwargs={ - "timestamp": self.timestamp, - "user_name": self.user_name, - "event_id": self.event_id, - "user_id": self.user_id, - "room_name": self.room_name, - "room_id": self.room_id, - } - ) + if self.role == "human": + return HumanMessage( + content=self.message, + role=self.user_name, # "chat" + additional_kwargs={ + "timestamp": self.timestamp, + "user_name": self.user_name, + "event_id": self.event_id, + "user_id": self.user_id, + "room_name": self.room_name, + "room_id": self.room_id, + } + ) + elif self.role == "ai": + return AIMessage( + content=self.message, + role=self.user_name, # "chat" + additional_kwargs={ + "timestamp": self.timestamp, + "user_name": self.user_name, + "event_id": self.event_id, + "user_id": self.user_id, + "room_name": self.room_name, + "room_id": self.room_id, + } + ) + elif self.role == "system": + return SystemMessage( + content=self.message, + role=self.user_name, # "chat" + additional_kwargs={ + "timestamp": self.timestamp, + "user_name": self.user_name, + "event_id": self.event_id, + "user_id": self.user_id, + "room_name": self.room_name, + "room_id": self.room_id, + } + ) + else: + return ChatMessage( + content=self.message, + role=self.user_name, # "chat" + additional_kwargs={ + "timestamp": self.timestamp, + "user_name": self.user_name, + "event_id": self.event_id, + "user_id": self.user_id, + "room_name": self.room_name, + "room_id": self.room_id, + } + ) def is_from(self, user_id): return self.user_id == user_id diff --git a/matrix_pygmalion_bot/bot/wrappers/langchain_koboldcpp.py b/matrix_pygmalion_bot/bot/wrappers/langchain_koboldcpp.py index 8a5ac53..7b42b6c 100644 --- a/matrix_pygmalion_bot/bot/wrappers/langchain_koboldcpp.py +++ b/matrix_pygmalion_bot/bot/wrappers/langchain_koboldcpp.py @@ -10,7 +10,7 @@ import functools from langchain.llms.base import LLM -from langchain.schema import BaseMessage +from langchain.schema import BaseMessage, AIMessage, HumanMessage, SystemMessage, ChatMessage logger = logging.getLogger(__name__) @@ -164,8 +164,9 @@ class KoboldCpp(LLM): for message in messages_dict: num_tokens += tokens_per_message for key, value in message.items(): - num_tokens += len(self.get_num_tokens(value)) - if key == "name": + if key == "content": + num_tokens += self.get_num_tokens(value) + elif key == "name": num_tokens += tokens_per_name num_tokens += 3 return num_tokens