chat_memory.py

from __future__ import annotations

from dataclasses import dataclass, field

from config import MAX_CONTEXT_CHARS, MAX_CONTEXT_MESSAGES, SYSTEM_PROMPT


Message = dict[str, str]


@dataclass
class ChatMemory:
    system_prompt: str = SYSTEM_PROMPT
    max_context_messages: int = MAX_CONTEXT_MESSAGES
    max_context_chars: int = MAX_CONTEXT_CHARS
    messages: list[Message] = field(default_factory=list)

    def __post_init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        self.messages = [{"role": "system", "content": self.system_prompt}]

    def add_user_message(self, content: str) -> None:
        self.messages.append({"role": "user", "content": content})

    def add_assistant_message(self, content: str) -> None:
        self.messages.append({"role": "assistant", "content": content})

    def get_messages(self) -> list[Message]:
        system_message = self.messages[0]
        recent_messages = self.messages[1:][-self.max_context_messages :]
        available_chars = max(0, self.max_context_chars - len(system_message["content"]))

        selected: list[Message] = []
        used_chars = len(system_message["content"])
        for message in reversed(recent_messages):
            remaining_chars = self.max_context_chars - used_chars
            if remaining_chars <= 0:
                break

            content = message["content"]
            message_chars = len(content)
            if message_chars > remaining_chars:
                if not selected:
                    selected.append(
                        {
                            "role": message["role"],
                            "content": _truncate_from_end(content, max(0, available_chars)),
                        }
                    )
                break

            selected.append(message.copy())
            used_chars += message_chars

        selected.reverse()
        return [system_message.copy(), *selected]


def _truncate_from_end(content: str, max_chars: int) -> str:
    if len(content) <= max_chars:
        return content
    if max_chars <= 80:
        return content[-max_chars:] if max_chars > 0 else ""

    marker = "[Earlier part truncated to fit context window]\n"
    keep_chars = max(0, max_chars - len(marker))
    return marker + content[-keep_chars:]