tayyab-chatbot / src /conversation.py
tayyab-077's picture
deploy chatbot
0885208
# src/conversation.py
from collections import deque
from typing import List, Dict, Optional
class ConversationMemory:
"""
Role-aware conversation memory stored as a deque of dicts:
{"role": "user"|"assistant"|"system", "content": "..."}
Default max_len keeps the most recent messages.
"""
def __init__(self, max_len: int = 60):
self.max_len = max_len
self.history = deque(maxlen=max_len)
def add(self, user: str, bot: str):
"""Add user message and assistant reply as two messages."""
if user is not None:
self.history.append({"role": "user", "content": user})
if bot is not None:
self.history.append({"role": "assistant", "content": bot})
def add_message(self, role: str, content: str):
if role not in ("user", "assistant", "system"):
raise ValueError("role must be 'user', 'assistant' or 'system'")
self.history.append({"role": role, "content": content})
def get_formatted(self, separator: str = "\n") -> str:
out_lines = []
for msg in self.history:
role_label = "System" if msg["role"] == "system" else ("User" if msg["role"] == "user" else "Assistant")
out_lines.append(f"{role_label}: {msg['content']}")
return separator.join(out_lines) + (separator if out_lines else "")
def clear(self):
self.history.clear()
def as_list(self) -> List[Dict[str, str]]:
return list(self.history)
def trim_to_recent_pairs(self, max_pairs: int = 12):
"""
Keep roughly the last `max_pairs` (user+assistant) pairs.
"""
items = list(self.history)
keep = items[-(2 * max_pairs):] if len(items) > 2 * max_pairs else items
self.history = deque(keep, maxlen=self.max_len)