Spaces:
Sleeping
Sleeping
File size: 2,707 Bytes
af8b301 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
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"]
@dataclass
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
@property
def history(self) -> List[Message]:
return list(self._history)
@property
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})"
)
|