| """Long-context baseline: the most recent visits that fit the token budget, oldest dropped first. |
| |
| The budget is counted with the configured tokenizer (`answer.tokenizer` in backbone.yaml). |
| """ |
| from __future__ import annotations |
|
|
| import time |
|
|
| from .base import BaseAgent, MemoryBuildResult, MemoryContext |
|
|
|
|
| class LongContextAgent(BaseAgent): |
| def __init__(self, *, max_context_tokens: int, margin: int, **params): |
| """max_context_tokens is the input budget in tokens; margin is held back from it for the prompt and question.""" |
| super().__init__(**params) |
| self.limit = max_context_tokens - margin |
| self._visits: list[tuple[int, str, int]] = [] |
|
|
| def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: |
| t0 = time.time() |
| self._visits.append((visit_number, text, self.count_tokens(text))) |
| return MemoryBuildResult(success=True, method="long_context", action="append_visit", |
| visit_number=visit_number, chunk_count=len(self._visits), |
| added_chunks=1, time_cost=time.time() - t0) |
|
|
| def build_memory_context(self, question: str, as_of: str = "", |
| display_question: str | None = None) -> MemoryContext: |
| total = sum(n for _, _, n in self._visits) |
| dropped: set[int] = set() |
| i = 0 |
| while total > self.limit and i < len(self._visits): |
| dropped.add(self._visits[i][0]) |
| total -= self._visits[i][2] |
| i += 1 |
| mem = "\n\n".join(t for _, t, _ in self._visits[i:]) |
| return MemoryContext( |
| mem=mem, n_tok=total, |
| retrieved_visits=[v for v, _, _ in self._visits[i:]], |
| context_truncated=bool(dropped), |
| extra={"visits_in_mem": len(self._visits) - i, "visits_dropped": i}, |
| ) |
|
|
