backstory/src/utils/conversation.py

41 lines
1.5 KiB
Python

from pydantic import BaseModel, Field, PrivateAttr # type: ignore
from typing import List
from .message import Message
class Conversation(BaseModel):
Conversation_messages: List[Message] = Field(default=[], alias="messages")
def __len__(self):
return len(self.Conversation_messages)
def __iter__(self):
return iter(self.Conversation_messages)
def reset(self):
self.Conversation_messages = []
@property
def messages(self):
"""Return a copy of messages to prevent modification of the internal list."""
raise AttributeError("Cannot directly get messages. Use Conversation.add() or .reset()")
@messages.setter
def messages(self, value):
"""Control how messages can be set, or prevent setting altogether."""
raise AttributeError("Cannot directly set messages. Use Conversation.add() or .reset()")
def add(self, message: Message | List[Message]) -> None:
"""Add a Message(s) to the conversation."""
if isinstance(message, Message):
self.Conversation_messages.append(message)
else:
self.Conversation_messages.extend(message)
def get_summary(self) -> str:
"""Return a summary of the conversation."""
if not self.Conversation_messages:
return "Conversation is empty."
summary = f"Conversation:\n"
for i, message in enumerate(self.Conversation_messages, 1):
summary += f"\nMessage {i}:\n{message.get_summary()}\n"
return summary