forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from abc import ABC
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from langchain.memory.utils import get_prompt_input_key
|
|
from langchain.schema import AIMessage, BaseMemory, BaseMessage, HumanMessage
|
|
|
|
|
|
class ChatMessageHistory(BaseModel):
|
|
messages: List[BaseMessage] = Field(default_factory=list)
|
|
|
|
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 clear(self) -> None:
|
|
self.messages = []
|
|
|
|
|
|
class BaseChatMemory(BaseMemory, ABC):
|
|
chat_memory: ChatMessageHistory = Field(default_factory=ChatMessageHistory)
|
|
output_key: Optional[str] = None
|
|
input_key: Optional[str] = None
|
|
return_messages: bool = False
|
|
|
|
def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
|
|
"""Save context from this conversation to buffer."""
|
|
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
|
|
self.chat_memory.add_user_message(inputs[prompt_input_key])
|
|
self.chat_memory.add_ai_message(outputs[output_key])
|
|
|
|
def clear(self) -> None:
|
|
"""Clear memory contents."""
|
|
self.chat_memory.clear()
|