File size: 1,865 Bytes
750b88d | 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 | """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},
)
|