Spaces:
Sleeping
Sleeping
| """ | |
| Short-term conversational memory manager. | |
| Maintains a rolling buffer of the last N exchanges for both assistants. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| from dataclasses import dataclass, field | |
| from typing import List, Literal, Optional | |
| Role = Literal["user", "assistant", "system"] | |
| class Message: | |
| role: Role | |
| content: str | |
| timestamp: float = field(default_factory=time.time) | |
| def to_dict(self) -> dict: | |
| return {"role": self.role, "content": self.content} | |
| class ConversationMemory: | |
| """ | |
| Rolling buffer memory that keeps the last `max_exchanges` user/assistant pairs. | |
| A system prompt is always prepended when building the message list. | |
| """ | |
| def __init__(self, max_exchanges: int = 8, system_prompt: Optional[str] = None): | |
| self.max_exchanges = max_exchanges | |
| self.system_prompt: Optional[str] = system_prompt | |
| self._history: List[Message] = [] | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def add_user_message(self, content: str) -> None: | |
| self._history.append(Message(role="user", content=content)) | |
| self._trim() | |
| def add_assistant_message(self, content: str) -> None: | |
| self._history.append(Message(role="assistant", content=content)) | |
| self._trim() | |
| def get_messages(self) -> List[dict]: | |
| """Return the full message list ready for an LLM API call.""" | |
| messages: List[dict] = [] | |
| if self.system_prompt: | |
| messages.append({"role": "system", "content": self.system_prompt}) | |
| messages.extend(msg.to_dict() for msg in self._history) | |
| return messages | |
| def clear(self) -> None: | |
| self._history.clear() | |
| def set_system_prompt(self, prompt: str) -> None: | |
| self.system_prompt = prompt | |
| def history(self) -> List[Message]: | |
| return list(self._history) | |
| def turn_count(self) -> int: | |
| return sum(1 for m in self._history if m.role == "user") | |
| # ------------------------------------------------------------------ | |
| # Internal helpers | |
| # ------------------------------------------------------------------ | |
| def _trim(self) -> None: | |
| """Keep only the last max_exchanges * 2 messages (user + assistant pairs).""" | |
| max_messages = self.max_exchanges * 2 | |
| if len(self._history) > max_messages: | |
| self._history = self._history[-max_messages:] | |
| def __repr__(self) -> str: | |
| return ( | |
| f"ConversationMemory(turns={self.turn_count}, " | |
| f"max_exchanges={self.max_exchanges})" | |
| ) | |