Coverage for agentos/memory/conversation.py: 42%
146 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1""" # noqa: E501
2Conversation Memory with sliding window management.
4Manages multi-turn conversations with configurable window strategies:
5- Sliding window (FIFO with max turns)
6- Token-aware window (trim by token count)
7- Importance-weighted (keep high-importance turns, evict low)
8- Hybrid (combine token budget + importance scoring)
9"""
11from __future__ import annotations
13from dataclasses import dataclass, field
14from enum import Enum
15from typing import Any
18class WindowStrategy(Enum):
19 SLIDING = "sliding"
20 """FIFO: keep last N turns, evict oldest."""
22 TOKEN_AWARE = "token_aware"
23 """Keep as many turns as fit within token budget."""
25 IMPORTANCE = "importance"
26 """Keep high-importance turns, evict lowest scores."""
28 HYBRID = "hybrid"
29 """Token budget + importance scoring combined."""
32@dataclass
33class ConversationTurn:
34 """Single turn in a conversation."""
36 role: str
37 """'user', 'assistant', 'system', 'tool'."""
39 content: str
40 timestamp: float = 0.0
41 token_count: int = 0
42 importance: float = 0.5
43 """0.0 = least important, 1.0 = most important."""
45 metadata: dict[str, Any] = field(default_factory=dict)
48@dataclass
49class WindowConfig:
50 """Configuration for conversation window management."""
52 strategy: WindowStrategy = WindowStrategy.SLIDING
54 max_turns: int = 20
55 """Max conversation turns (sliding window)."""
57 max_tokens: int = 8000
58 """Max total token budget (token_aware / hybrid)."""
60 importance_threshold: float = 0.3
61 """Minimum importance score to keep (importance / hybrid)."""
63 system_prompt: str | None = None
64 """System prompt always kept at top of window."""
66 preserve_last_n: int = 2
67 """Always keep the last N turns regardless of eviction rules."""
70class ConversationMemory:
71 """
72 Multi-turn conversation memory with sliding window strategies.
74 Example::
76 mem = ConversationMemory(WindowConfig(strategy=WindowStrategy.HYBRID, max_tokens=4000))
77 mem.add_turn(ConversationTurn(role="user", content="Hello"))
78 mem.add_turn(ConversationTurn(role="assistant", content="Hi! How can I help?"))
79 messages = mem.get_messages() # [{"role": "user", "content": "Hello"}, ...]
80 """
82 def __init__(self, config: WindowConfig | None = None):
83 self.config = config or WindowConfig()
84 self._turns: list[ConversationTurn] = []
85 self._token_count_cache: int = 0
87 def add_turn(self, turn: ConversationTurn) -> None:
88 """Add a turn and apply window eviction if needed."""
89 self._turns.append(turn)
90 self._token_count_cache += (
91 turn.token_count if turn.token_count > 0 else self._estimate_tokens(turn.content)
92 )
93 self._apply_window()
95 def add_user_message(self, content: str, importance: float = 0.5) -> None:
96 self.add_turn(
97 ConversationTurn(
98 role="user",
99 content=content,
100 importance=importance,
101 token_count=self._estimate_tokens(content),
102 )
103 )
105 def add_assistant_message(self, content: str, importance: float = 0.5) -> None:
106 self.add_turn(
107 ConversationTurn(
108 role="assistant",
109 content=content,
110 importance=importance,
111 token_count=self._estimate_tokens(content),
112 )
113 )
115 def add_system_message(self, content: str) -> None:
116 self.add_turn(
117 ConversationTurn(
118 role="system",
119 content=content,
120 importance=1.0,
121 token_count=self._estimate_tokens(content),
122 )
123 )
125 def _apply_window(self) -> None:
126 """Apply the configured window strategy to evict excess turns."""
127 strategy = self.config.strategy
129 if strategy == WindowStrategy.SLIDING:
130 self._evict_sliding()
131 elif strategy == WindowStrategy.TOKEN_AWARE:
132 while (
133 self._token_count_cache > self.config.max_tokens
134 and len(self._turns) > self.config.preserve_last_n
135 ):
136 self._evict_one(0)
137 elif strategy == WindowStrategy.IMPORTANCE:
138 self._evict_by_importance()
139 elif strategy == WindowStrategy.HYBRID:
140 self._evict_hybrid()
142 def _evict_sliding(self) -> None:
143 """FIFO: remove oldest turns exceeding max_turns."""
144 preserve = self.config.preserve_last_n
145 max_keep = self.config.max_turns
147 while len(self._turns) > max_keep:
148 evict_idx = 0
149 # Don't evict system prompt
150 if self._turns[0].role == "system":
151 evict_idx = 1
152 # Don't evict preserved last N turns
153 if len(self._turns) - evict_idx <= preserve:
154 break
155 self._evict_one(evict_idx)
157 def _evict_by_importance(self) -> None:
158 """Evict lowest-importance turns above threshold."""
159 preserve = self.config.preserve_last_n
160 threshold = self.config.importance_threshold
162 while True:
163 candidates = [
164 (i, t)
165 for i, t in enumerate(self._turns)
166 if t.role != "system"
167 and i < len(self._turns) - preserve
168 and t.importance < threshold
169 ]
170 if not candidates:
171 break
173 # Evict the least important
174 idx, _ = min(candidates, key=lambda x: x[1].importance)
175 self._evict_one(idx)
176 if not any(
177 t.importance < threshold
178 for i, t in enumerate(self._turns)
179 if t.role != "system" and i < len(self._turns) - preserve
180 ):
181 break
183 def _evict_hybrid(self) -> None:
184 """Token budget + importance scoring combined."""
185 preserve = self.config.preserve_last_n
186 threshold = self.config.importance_threshold
188 # First, evict low-importance turns within budget
189 while self._token_count_cache > self.config.max_tokens:
190 candidates = [
191 (i, t)
192 for i, t in enumerate(self._turns)
193 if t.role != "system"
194 and i < len(self._turns) - preserve
195 and t.importance < threshold
196 ]
197 if not candidates:
198 # Fall back to evicting oldest non-system turn
199 oldest_idx = -1
200 for i, t in enumerate(self._turns):
201 if t.role != "system" and i < len(self._turns) - preserve:
202 oldest_idx = i
203 break
204 if oldest_idx == -1:
205 break
206 self._evict_one(oldest_idx)
207 else:
208 idx, _ = min(candidates, key=lambda x: x[1].importance)
209 self._evict_one(idx)
211 def _evict_one(self, index: int) -> None:
212 """Remove a single turn at given index."""
213 if 0 <= index < len(self._turns):
214 turn = self._turns.pop(index)
215 self._token_count_cache -= (
216 turn.token_count if turn.token_count > 0 else self._estimate_tokens(turn.content)
217 )
218 self._token_count_cache = max(0, self._token_count_cache)
220 def get_messages(self) -> list[dict[str, str]]:
221 """Return conversation as list of dicts (OpenAI chat format)."""
222 msgs: list[dict[str, str]] = []
223 if self.config.system_prompt:
224 msgs.append({"role": "system", "content": self.config.system_prompt})
225 for turn in self._turns:
226 msgs.append({"role": turn.role, "content": turn.content})
227 return msgs
229 def get_turns(self) -> list[ConversationTurn]:
230 return list(self._turns)
232 @property
233 def turn_count(self) -> int:
234 return len(self._turns)
236 @property
237 def token_count(self) -> int:
238 return self._token_count_cache
240 def clear(self) -> None:
241 """Reset conversation memory."""
242 self._turns.clear()
243 self._token_count_cache = 0
245 def to_summary(self) -> str:
246 """Generate a brief summary of the conversation memory."""
247 turns = self._turns
248 if not turns:
249 return "Empty conversation."
251 lines = [
252 f"Total turns: {len(turns)}",
253 f"Total tokens (est.): {self._token_count_cache}",
254 f"First turn: [{turns[0].role}] {turns[0].content[:80]}...",
255 ]
256 if len(turns) > 1:
257 lines.append(f"Last turn: [{turns[-1].role}] {turns[-1].content[:80]}...")
258 return "\n".join(lines)
260 @staticmethod
261 def _estimate_tokens(text: str) -> int:
262 """Rough token estimation: ~4 chars per token."""
263 return max(1, len(text) // 4)
265 def __len__(self) -> int:
266 return len(self._turns)
268 def __repr__(self) -> str:
269 return f"ConversationMemory(turns={len(self._turns)}, tokens={self._token_count_cache}, strategy={self.config.strategy.value})" # noqa: E501
271 # ── Persistence (v1.14.9) ────────────────
273 def get_state(self) -> dict[str, Any]:
274 """Export conversation memory state for persistence."""
275 return {
276 "config": {
277 "strategy": self.config.strategy.value,
278 "max_turns": self.config.max_turns,
279 "max_tokens": self.config.max_tokens,
280 "importance_threshold": self.config.importance_threshold,
281 "system_prompt": self.config.system_prompt,
282 "preserve_last_n": self.config.preserve_last_n,
283 },
284 "turns": [
285 {
286 "role": turn.role,
287 "content": turn.content,
288 "timestamp": turn.timestamp,
289 "token_count": turn.token_count,
290 "importance": turn.importance,
291 "metadata": turn.metadata,
292 }
293 for turn in self._turns
294 ],
295 "token_count_cache": self._token_count_cache,
296 }
298 def restore_state(self, state: dict[str, Any]) -> None:
299 """Restore conversation memory from a persisted snapshot."""
300 config_data = state.get("config", {})
301 self.config = WindowConfig(
302 strategy=WindowStrategy(config_data.get("strategy", "sliding")),
303 max_turns=config_data.get("max_turns", 20),
304 max_tokens=config_data.get("max_tokens", 8000),
305 importance_threshold=config_data.get("importance_threshold", 0.3),
306 system_prompt=config_data.get("system_prompt"),
307 preserve_last_n=config_data.get("preserve_last_n", 2),
308 )
309 self._turns = []
310 for turn_data in state.get("turns", []):
311 self._turns.append(
312 ConversationTurn(
313 role=turn_data.get("role", "user"),
314 content=turn_data.get("content", ""),
315 timestamp=turn_data.get("timestamp", 0.0),
316 token_count=turn_data.get("token_count", 0),
317 importance=turn_data.get("importance", 0.5),
318 metadata=turn_data.get("metadata", {}),
319 )
320 )
321 self._token_count_cache = state.get("token_count_cache", 0)