« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Token budget allocation for working memory assembly."""
3from __future__ import annotations
5from lexigram.ai.memory.config import WorkingMemoryConfig
8class TokenBudgetAllocator:
9 """Distributes a total token budget across working memory sections.
11 The budget is split in this order:
12 1. System prompt receives a fixed allocation.
13 2. The remaining budget is divided among recent turns, episodic recall,
14 semantic facts, and tool descriptions using the configured fractions.
15 """
17 def __init__(self, config: WorkingMemoryConfig | None = None) -> None:
18 """Initialise with optional config.
20 Args:
21 config: Working memory configuration; uses defaults if None.
22 """
23 self._config = config or WorkingMemoryConfig()
25 def allocate(self, total_tokens: int) -> dict[str, int]:
26 """Compute token allocations for each memory section.
28 Args:
29 total_tokens: Total token budget available.
31 Returns:
32 Mapping of section name to token allocation.
33 """
34 system = min(self._config.system_prompt_tokens, total_tokens)
35 remaining = max(0, total_tokens - system)
37 sections = {
38 "recent_turns": int(remaining * self._config.recent_turns_fraction),
39 "episodic": int(remaining * self._config.episodic_fraction),
40 "semantic": int(remaining * self._config.semantic_fraction),
41 "tool_descriptions": int(
42 remaining * self._config.tool_descriptions_fraction
43 ),
44 }
45 # Distribute any rounding remainder to the largest bucket
46 allocated = sum(sections.values())
47 remainder = remaining - allocated
48 if remainder > 0 and sections:
49 largest = max(sections, key=sections.__getitem__)
50 sections[largest] += remainder
52 return {"system_prompt": system, **sections}
54 def budget_for(self, section: str, total_tokens: int) -> int:
55 """Return the token budget for a single named section.
57 Args:
58 section: Section name (e.g. 'episodic', 'semantic').
59 total_tokens: Total token budget.
61 Returns:
62 Token count allocated to the requested section.
63 """
64 return self.allocate(total_tokens).get(section, 0)
67__all__ = ["TokenBudgetAllocator"]