Rename embedchain to mem0 and open sourcing code for long term memory (#1474)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Taranjeet Singh
co-authored by
Deshraj Yadav
Deshraj Yadav
parent
83e8c97295
commit
f842a92e25
@@ -0,0 +1,127 @@
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.core.db.database import get_session
|
||||
from embedchain.core.db.models import ChatHistory as ChatHistoryModel
|
||||
from embedchain.memory.message import ChatMessage
|
||||
from embedchain.memory.utils import merge_metadata_dict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChatHistory:
|
||||
def __init__(self) -> None:
|
||||
self.db_session = get_session()
|
||||
|
||||
def add(self, app_id, session_id, chat_message: ChatMessage) -> Optional[str]:
|
||||
memory_id = str(uuid.uuid4())
|
||||
metadata_dict = merge_metadata_dict(chat_message.human_message.metadata, chat_message.ai_message.metadata)
|
||||
if metadata_dict:
|
||||
metadata = self._serialize_json(metadata_dict)
|
||||
self.db_session.add(
|
||||
ChatHistoryModel(
|
||||
app_id=app_id,
|
||||
id=memory_id,
|
||||
session_id=session_id,
|
||||
question=chat_message.human_message.content,
|
||||
answer=chat_message.ai_message.content,
|
||||
metadata=metadata if metadata_dict else "{}",
|
||||
)
|
||||
)
|
||||
try:
|
||||
self.db_session.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"Error adding chat memory to db: {e}")
|
||||
self.db_session.rollback()
|
||||
return None
|
||||
|
||||
logger.info(f"Added chat memory to db with id: {memory_id}")
|
||||
return memory_id
|
||||
|
||||
def delete(self, app_id: str, session_id: Optional[str] = None):
|
||||
"""
|
||||
Delete all chat history for a given app_id and session_id.
|
||||
This is useful for deleting chat history for a given user.
|
||||
|
||||
:param app_id: The app_id to delete chat history for
|
||||
:param session_id: The session_id to delete chat history for
|
||||
|
||||
:return: None
|
||||
"""
|
||||
params = {"app_id": app_id}
|
||||
if session_id:
|
||||
params["session_id"] = session_id
|
||||
self.db_session.query(ChatHistoryModel).filter_by(**params).delete()
|
||||
try:
|
||||
self.db_session.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting chat history: {e}")
|
||||
self.db_session.rollback()
|
||||
|
||||
def get(
|
||||
self, app_id, session_id: str = "default", num_rounds=10, fetch_all: bool = False, display_format=False
|
||||
) -> list[ChatMessage]:
|
||||
"""
|
||||
Get the chat history for a given app_id.
|
||||
|
||||
param: app_id - The app_id to get chat history
|
||||
param: session_id (optional) - The session_id to get chat history. Defaults to "default"
|
||||
param: num_rounds (optional) - The number of rounds to get chat history. Defaults to 10
|
||||
param: fetch_all (optional) - Whether to fetch all chat history or not. Defaults to False
|
||||
param: display_format (optional) - Whether to return the chat history in display format. Defaults to False
|
||||
"""
|
||||
params = {"app_id": app_id}
|
||||
if not fetch_all:
|
||||
params["session_id"] = session_id
|
||||
results = (
|
||||
self.db_session.query(ChatHistoryModel).filter_by(**params).order_by(ChatHistoryModel.created_at.asc())
|
||||
)
|
||||
results = results.limit(num_rounds) if not fetch_all else results
|
||||
history = []
|
||||
for result in results:
|
||||
metadata = self._deserialize_json(metadata=result.meta_data or "{}")
|
||||
# Return list of dict if display_format is True
|
||||
if display_format:
|
||||
history.append(
|
||||
{
|
||||
"session_id": result.session_id,
|
||||
"human": result.question,
|
||||
"ai": result.answer,
|
||||
"metadata": result.meta_data,
|
||||
"timestamp": result.created_at,
|
||||
}
|
||||
)
|
||||
else:
|
||||
memory = ChatMessage()
|
||||
memory.add_user_message(result.question, metadata=metadata)
|
||||
memory.add_ai_message(result.answer, metadata=metadata)
|
||||
history.append(memory)
|
||||
return history
|
||||
|
||||
def count(self, app_id: str, session_id: Optional[str] = None):
|
||||
"""
|
||||
Count the number of chat messages for a given app_id and session_id.
|
||||
|
||||
:param app_id: The app_id to count chat history for
|
||||
:param session_id: The session_id to count chat history for
|
||||
|
||||
:return: The number of chat messages for a given app_id and session_id
|
||||
"""
|
||||
# Rewrite the logic below with sqlalchemy
|
||||
params = {"app_id": app_id}
|
||||
if session_id:
|
||||
params["session_id"] = session_id
|
||||
return self.db_session.query(ChatHistoryModel).filter_by(**params).count()
|
||||
|
||||
@staticmethod
|
||||
def _serialize_json(metadata: dict[str, Any]):
|
||||
return json.dumps(metadata)
|
||||
|
||||
@staticmethod
|
||||
def _deserialize_json(metadata: str):
|
||||
return json.loads(metadata)
|
||||
|
||||
def close_connection(self):
|
||||
self.connection.close()
|
||||
@@ -0,0 +1,74 @@
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.helpers.json_serializable import JSONSerializable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseMessage(JSONSerializable):
|
||||
"""
|
||||
The base abstract message class.
|
||||
|
||||
Messages are the inputs and outputs of Models.
|
||||
"""
|
||||
|
||||
# The string content of the message.
|
||||
content: str
|
||||
|
||||
# The created_by of the message. AI, Human, Bot etc.
|
||||
created_by: str
|
||||
|
||||
# Any additional info.
|
||||
metadata: dict[str, Any]
|
||||
|
||||
def __init__(self, content: str, created_by: str, metadata: Optional[dict[str, Any]] = None) -> None:
|
||||
super().__init__()
|
||||
self.content = content
|
||||
self.created_by = created_by
|
||||
self.metadata = metadata
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
"""Type of the Message, used for serialization."""
|
||||
|
||||
@classmethod
|
||||
def is_lc_serializable(cls) -> bool:
|
||||
"""Return whether this class is serializable."""
|
||||
return True
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.created_by}: {self.content}"
|
||||
|
||||
|
||||
class ChatMessage(JSONSerializable):
|
||||
"""
|
||||
The base abstract chat message class.
|
||||
|
||||
Chat messages are the pair of (question, answer) conversation
|
||||
between human and model.
|
||||
"""
|
||||
|
||||
human_message: Optional[BaseMessage] = None
|
||||
ai_message: Optional[BaseMessage] = None
|
||||
|
||||
def add_user_message(self, message: str, metadata: Optional[dict] = None):
|
||||
if self.human_message:
|
||||
logger.info(
|
||||
"Human message already exists in the chat message,\
|
||||
overwriting it with new message."
|
||||
)
|
||||
|
||||
self.human_message = BaseMessage(content=message, created_by="human", metadata=metadata)
|
||||
|
||||
def add_ai_message(self, message: str, metadata: Optional[dict] = None):
|
||||
if self.ai_message:
|
||||
logger.info(
|
||||
"AI message already exists in the chat message,\
|
||||
overwriting it with new message."
|
||||
)
|
||||
|
||||
self.ai_message = BaseMessage(content=message, created_by="ai", metadata=metadata)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.human_message}\n{self.ai_message}"
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
def merge_metadata_dict(left: Optional[dict[str, Any]], right: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:
|
||||
"""
|
||||
Merge the metadatas of two BaseMessage types.
|
||||
|
||||
Args:
|
||||
left (dict[str, Any]): metadata of human message
|
||||
right (dict[str, Any]): metadata of AI message
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: combined metadata dict with dedup
|
||||
to be saved in db.
|
||||
"""
|
||||
if not left and not right:
|
||||
return None
|
||||
elif not left:
|
||||
return right
|
||||
elif not right:
|
||||
return left
|
||||
|
||||
merged = left.copy()
|
||||
for k, v in right.items():
|
||||
if k not in merged:
|
||||
merged[k] = v
|
||||
elif type(merged[k]) != type(v):
|
||||
raise ValueError(f'additional_kwargs["{k}"] already exists in this message,' " but with a different type.")
|
||||
elif isinstance(merged[k], str):
|
||||
merged[k] += v
|
||||
elif isinstance(merged[k], dict):
|
||||
merged[k] = merge_metadata_dict(merged[k], v)
|
||||
else:
|
||||
raise ValueError(f"Additional kwargs key {k} already exists in this message.")
|
||||
return merged
|
||||
Reference in New Issue
Block a user