Spaces:
Sleeping
Sleeping
| # 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) | |