src/memory.py

"""Conversation memory: keeps recent turns + tracked entities so follow-ups
like "now compare that with last year's report" resolve without re-stating
the company or re-uploading files.
"""

from __future__ import annotations

from dataclasses import dataclass, field

MAX_TURNS = 8


@dataclass
class Turn:
    question: str
    answer: str
    route: str


@dataclass
class ConversationMemory:
    turns: list[Turn] = field(default_factory=list)
    active_documents: list[str] = field(default_factory=list)

    def add(self, question: str, answer: str, route: str):
        self.turns.append(Turn(question, answer, route))
        self.turns = self.turns[-MAX_TURNS:]

    def set_documents(self, doc_ids: list[str]):
        self.active_documents = doc_ids

    def history_text(self, max_chars: int = 4000) -> str:
        """Compact transcript passed to agents for coreference resolution."""
        parts = []
        for t in self.turns:
            answer = t.answer if len(t.answer) < 600 else t.answer[:600] + " …"
            parts.append(f"User: {t.question}\nAssistant ({t.route}): {answer}")
        text = "\n\n".join(parts)
        return text[-max_chars:]

    def clear(self):
        self.turns.clear()