File size: 25,947 Bytes
23804b3 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 |
"""
Working Memory Management System with Attention-based Focus and Context Switching
Implements dynamic attention mechanisms and context management for cognitive agents
"""
import sqlite3
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, asdict
import logging
from pathlib import Path
import heapq
import threading
import time
logger = logging.getLogger(__name__)
@dataclass
class WorkingMemoryItem:
"""Individual item in working memory"""
id: str
content: str
item_type: str # goal, observation, hypothesis, plan, etc.
priority: float # 0.0-1.0, higher is more important
activation_level: float # 0.0-1.0, current activation
created_at: datetime
last_accessed: datetime
access_count: int
decay_rate: float # how quickly activation decays
context_tags: List[str]
source_agent: str
related_items: List[str] # IDs of related items
@dataclass
class AttentionFocus:
"""Current attention focus with weighted priorities"""
focus_id: str
focus_type: str # task, threat, goal, etc.
focus_items: List[str] # Working memory item IDs
attention_weight: float # 0.0-1.0
duration: timedelta
created_at: datetime
metadata: Dict[str, Any]
class WorkingMemoryManager:
"""Advanced working memory with attention-based focus management"""
def __init__(self, db_path: str = "data/cognitive/working_memory.db",
capacity: int = 50, decay_interval: float = 30.0):
"""Initialize working memory system"""
self.db_path = Path(db_path)
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self.capacity = capacity # Maximum items in working memory
self.decay_interval = decay_interval # Seconds between decay updates
self._init_database()
self._memory_items = {} # In-memory cache
self._attention_focus = None
self._attention_history = []
# Start background decay process
self._decay_thread = threading.Thread(target=self._decay_loop, daemon=True)
self._decay_running = True
self._decay_thread.start()
# Load existing items
self._load_working_memory()
def _init_database(self):
"""Initialize database schemas"""
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS working_memory_items (
id TEXT PRIMARY KEY,
content TEXT NOT NULL,
item_type TEXT NOT NULL,
priority REAL NOT NULL,
activation_level REAL NOT NULL,
created_at TEXT NOT NULL,
last_accessed TEXT NOT NULL,
access_count INTEGER DEFAULT 0,
decay_rate REAL DEFAULT 0.1,
context_tags TEXT,
source_agent TEXT,
related_items TEXT,
updated_at TEXT DEFAULT CURRENT_TIMESTAMP
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS attention_focus_log (
id TEXT PRIMARY KEY,
focus_type TEXT NOT NULL,
focus_items TEXT,
attention_weight REAL NOT NULL,
duration_seconds REAL,
created_at TEXT NOT NULL,
ended_at TEXT,
metadata TEXT,
agent_id TEXT
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS context_switches (
id TEXT PRIMARY KEY,
from_focus TEXT,
to_focus TEXT,
switch_reason TEXT,
switch_cost REAL,
timestamp TEXT DEFAULT CURRENT_TIMESTAMP,
agent_id TEXT
)
""")
# Create indices
conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_priority ON working_memory_items(priority DESC)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_activation ON working_memory_items(activation_level DESC)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_type ON working_memory_items(item_type)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_agent ON working_memory_items(source_agent)")
def add_item(self, content: str, item_type: str, priority: float = 0.5,
source_agent: str = "", context_tags: List[str] = None) -> str:
"""Add item to working memory with attention-based priority"""
try:
item_id = str(uuid.uuid4())
item = WorkingMemoryItem(
id=item_id,
content=content,
item_type=item_type,
priority=priority,
activation_level=priority, # Initial activation equals priority
created_at=datetime.now(),
last_accessed=datetime.now(),
access_count=0,
decay_rate=0.1, # Default decay rate
context_tags=context_tags or [],
source_agent=source_agent,
related_items=[]
)
# Check capacity and evict if necessary
if len(self._memory_items) >= self.capacity:
self._evict_lowest_activation()
# Store in memory and database
self._memory_items[item_id] = item
self._store_item_to_db(item)
# Update attention focus if this is high priority
if priority > 0.7 and (not self._attention_focus or
priority > self._attention_focus.attention_weight):
self._update_attention_focus(item_id, item_type, priority)
logger.info(f"Added working memory item: {item_type} (priority: {priority:.2f})")
return item_id
except Exception as e:
logger.error(f"Error adding working memory item: {e}")
return ""
def get_item(self, item_id: str) -> Optional[WorkingMemoryItem]:
"""Retrieve item from working memory and update activation"""
try:
if item_id in self._memory_items:
item = self._memory_items[item_id]
# Update access statistics
item.last_accessed = datetime.now()
item.access_count += 1
# Boost activation on access (but cap at 1.0)
activation_boost = min(0.2, 1.0 - item.activation_level)
item.activation_level = min(1.0, item.activation_level + activation_boost)
# Update in database
self._update_item_in_db(item)
logger.debug(f"Retrieved working memory item: {item_id[:8]}...")
return item
logger.warning(f"Working memory item not found: {item_id}")
return None
except Exception as e:
logger.error(f"Error retrieving working memory item: {e}")
return None
def get_active_items(self, min_activation: float = 0.3,
item_type: str = "", limit: int = 20) -> List[WorkingMemoryItem]:
"""Get currently active items above activation threshold"""
try:
active_items = []
for item in self._memory_items.values():
if (item.activation_level >= min_activation and
(not item_type or item.item_type == item_type)):
active_items.append(item)
# Sort by activation level (highest first)
active_items.sort(key=lambda x: x.activation_level, reverse=True)
logger.info(f"Retrieved {len(active_items[:limit])} active items")
return active_items[:limit]
except Exception as e:
logger.error(f"Error getting active items: {e}")
return []
def focus_attention(self, focus_type: str, item_ids: List[str],
attention_weight: float = 0.8, agent_id: str = "") -> str:
"""Focus attention on specific items"""
try:
focus_id = str(uuid.uuid4())
# End current focus if exists
if self._attention_focus:
self._end_attention_focus()
# Create new attention focus
new_focus = AttentionFocus(
focus_id=focus_id,
focus_type=focus_type,
focus_items=item_ids,
attention_weight=attention_weight,
duration=timedelta(0),
created_at=datetime.now(),
metadata={'agent_id': agent_id}
)
self._attention_focus = new_focus
# Boost activation of focused items
for item_id in item_ids:
if item_id in self._memory_items:
item = self._memory_items[item_id]
item.activation_level = min(1.0, item.activation_level + 0.3)
self._update_item_in_db(item)
# Store focus in database
self._store_attention_focus(new_focus)
logger.info(f"Focused attention on {focus_type}: {len(item_ids)} items")
return focus_id
except Exception as e:
logger.error(f"Error focusing attention: {e}")
return ""
def switch_context(self, new_focus_type: str, new_item_ids: List[str],
switch_reason: str = "", agent_id: str = "") -> Dict[str, Any]:
"""Switch attention context with cost calculation"""
try:
switch_result = {
'switch_id': str(uuid.uuid4()),
'from_focus': None,
'to_focus': new_focus_type,
'switch_cost': 0.0,
'success': False
}
# Calculate switch cost based on current focus
if self._attention_focus:
switch_result['from_focus'] = self._attention_focus.focus_type
# Cost factors:
# 1. How long we've been in current focus
current_duration = datetime.now() - self._attention_focus.created_at
duration_cost = min(current_duration.total_seconds() / 300.0, 0.3) # Max 5min
# 2. Number of active items being abandoned
abandoned_items = len(self._attention_focus.focus_items)
abandonment_cost = min(abandoned_items * 0.1, 0.4)
# 3. Similarity between old and new focus
similarity_discount = self._calculate_focus_similarity(
self._attention_focus.focus_items, new_item_ids
)
total_cost = duration_cost + abandonment_cost - similarity_discount
switch_result['switch_cost'] = max(0.0, min(total_cost, 1.0))
# Record context switch
self._record_context_switch(
self._attention_focus.focus_type,
new_focus_type,
switch_reason,
switch_result['switch_cost'],
agent_id
)
# Perform the switch
focus_id = self.focus_attention(new_focus_type, new_item_ids, agent_id=agent_id)
switch_result['success'] = bool(focus_id)
logger.info(f"Context switch: {switch_result['from_focus']} -> {new_focus_type} (cost: {switch_result['switch_cost']:.3f})")
return switch_result
except Exception as e:
logger.error(f"Error switching context: {e}")
return {'error': str(e)}
def get_current_focus(self) -> Optional[AttentionFocus]:
"""Get current attention focus"""
return self._attention_focus
def decay_memory(self):
"""Apply decay to all working memory items"""
try:
decayed_count = 0
evicted_items = []
for item_id, item in list(self._memory_items.items()):
# Apply decay based on time since last access
time_since_access = datetime.now() - item.last_accessed
decay_amount = item.decay_rate * (time_since_access.total_seconds() / 60.0)
item.activation_level = max(0.0, item.activation_level - decay_amount)
decayed_count += 1
# Evict items with very low activation
if item.activation_level < 0.05:
evicted_items.append(item_id)
else:
# Update in database
self._update_item_in_db(item)
# Remove evicted items
for item_id in evicted_items:
del self._memory_items[item_id]
self._remove_item_from_db(item_id)
if evicted_items:
logger.info(f"Memory decay: {decayed_count} items decayed, {len(evicted_items)} evicted")
except Exception as e:
logger.error(f"Error during memory decay: {e}")
def find_related_items(self, item_id: str, max_items: int = 5) -> List[WorkingMemoryItem]:
"""Find items related to the given item"""
try:
if item_id not in self._memory_items:
return []
source_item = self._memory_items[item_id]
related_items = []
for other_id, other_item in self._memory_items.items():
if other_id == item_id:
continue
# Calculate relatedness score
relatedness = 0.0
# Same type bonus
if source_item.item_type == other_item.item_type:
relatedness += 0.3
# Shared context tags
shared_tags = set(source_item.context_tags) & set(other_item.context_tags)
if shared_tags:
relatedness += len(shared_tags) * 0.2
# Same source agent
if source_item.source_agent == other_item.source_agent:
relatedness += 0.2
# Temporal proximity
time_diff = abs((source_item.created_at - other_item.created_at).total_seconds())
if time_diff < 300: # Within 5 minutes
relatedness += 0.3 * (300 - time_diff) / 300
if relatedness > 0.1: # Minimum relatedness threshold
related_items.append((other_item, relatedness))
# Sort by relatedness and return top items
related_items.sort(key=lambda x: x[1], reverse=True)
return [item for item, score in related_items[:max_items]]
except Exception as e:
logger.error(f"Error finding related items: {e}")
return []
def _update_attention_focus(self, item_id: str, item_type: str, priority: float):
"""Update current attention focus"""
if self._attention_focus:
self._end_attention_focus()
self.focus_attention(item_type, [item_id], priority)
def _end_attention_focus(self):
"""End current attention focus"""
if self._attention_focus:
# Update duration
self._attention_focus.duration = datetime.now() - self._attention_focus.created_at
# Add to history
self._attention_history.append(self._attention_focus)
# Update in database
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
UPDATE attention_focus_log SET
ended_at = ?,
duration_seconds = ?
WHERE id = ?
""", (
datetime.now().isoformat(),
self._attention_focus.duration.total_seconds(),
self._attention_focus.focus_id
))
self._attention_focus = None
def _evict_lowest_activation(self):
"""Evict item with lowest activation to make space"""
if not self._memory_items:
return
lowest_item_id = min(
self._memory_items.keys(),
key=lambda x: self._memory_items[x].activation_level
)
del self._memory_items[lowest_item_id]
self._remove_item_from_db(lowest_item_id)
logger.debug(f"Evicted working memory item: {lowest_item_id[:8]}...")
def _calculate_focus_similarity(self, items1: List[str], items2: List[str]) -> float:
"""Calculate similarity between two sets of focus items"""
if not items1 or not items2:
return 0.0
set1 = set(items1)
set2 = set(items2)
intersection = len(set1 & set2)
union = len(set1 | set2)
return intersection / union if union > 0 else 0.0
def _record_context_switch(self, from_focus: str, to_focus: str,
reason: str, cost: float, agent_id: str):
"""Record context switch in database"""
try:
switch_id = str(uuid.uuid4())
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
INSERT INTO context_switches (
id, from_focus, to_focus, switch_reason,
switch_cost, agent_id
) VALUES (?, ?, ?, ?, ?, ?)
""", (switch_id, from_focus, to_focus, reason, cost, agent_id))
except Exception as e:
logger.error(f"Error recording context switch: {e}")
def _decay_loop(self):
"""Background thread for memory decay"""
while self._decay_running:
try:
time.sleep(self.decay_interval)
self.decay_memory()
except Exception as e:
logger.error(f"Error in decay loop: {e}")
def _load_working_memory(self):
"""Load working memory items from database"""
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.execute("""
SELECT * FROM working_memory_items
ORDER BY activation_level DESC
LIMIT ?
""", (self.capacity,))
for row in cursor.fetchall():
item = WorkingMemoryItem(
id=row[0],
content=row[1],
item_type=row[2],
priority=row[3],
activation_level=row[4],
created_at=datetime.fromisoformat(row[5]),
last_accessed=datetime.fromisoformat(row[6]),
access_count=row[7],
decay_rate=row[8],
context_tags=json.loads(row[9]) if row[9] else [],
source_agent=row[10] or "",
related_items=json.loads(row[11]) if row[11] else []
)
self._memory_items[item.id] = item
logger.info(f"Loaded {len(self._memory_items)} working memory items")
except Exception as e:
logger.error(f"Error loading working memory: {e}")
def _store_item_to_db(self, item: WorkingMemoryItem):
"""Store item to database"""
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
INSERT INTO working_memory_items (
id, content, item_type, priority, activation_level,
created_at, last_accessed, access_count, decay_rate,
context_tags, source_agent, related_items
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""", (
item.id, item.content, item.item_type, item.priority,
item.activation_level, item.created_at.isoformat(),
item.last_accessed.isoformat(), item.access_count,
item.decay_rate, json.dumps(item.context_tags),
item.source_agent, json.dumps(item.related_items)
))
except Exception as e:
logger.error(f"Error storing item to database: {e}")
def _update_item_in_db(self, item: WorkingMemoryItem):
"""Update item in database"""
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
UPDATE working_memory_items SET
activation_level = ?, last_accessed = ?,
access_count = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ?
""", (
item.activation_level, item.last_accessed.isoformat(),
item.access_count, item.id
))
except Exception as e:
logger.error(f"Error updating item in database: {e}")
def _remove_item_from_db(self, item_id: str):
"""Remove item from database"""
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("DELETE FROM working_memory_items WHERE id = ?", (item_id,))
except Exception as e:
logger.error(f"Error removing item from database: {e}")
def _store_attention_focus(self, focus: AttentionFocus):
"""Store attention focus in database"""
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
INSERT INTO attention_focus_log (
id, focus_type, focus_items, attention_weight,
created_at, metadata, agent_id
) VALUES (?, ?, ?, ?, ?, ?, ?)
""", (
focus.focus_id, focus.focus_type,
json.dumps(focus.focus_items), focus.attention_weight,
focus.created_at.isoformat(), json.dumps(focus.metadata),
focus.metadata.get('agent_id', '')
))
except Exception as e:
logger.error(f"Error storing attention focus: {e}")
def get_working_memory_statistics(self) -> Dict[str, Any]:
"""Get comprehensive working memory statistics"""
try:
with sqlite3.connect(self.db_path) as conn:
stats = {
'current_capacity': len(self._memory_items),
'max_capacity': self.capacity,
'utilization': len(self._memory_items) / self.capacity
}
# Activation distribution
if self._memory_items:
activations = [item.activation_level for item in self._memory_items.values()]
stats['avg_activation'] = sum(activations) / len(activations)
stats['max_activation'] = max(activations)
stats['min_activation'] = min(activations)
# Item type distribution
type_counts = {}
for item in self._memory_items.values():
type_counts[item.item_type] = type_counts.get(item.item_type, 0) + 1
stats['item_types'] = type_counts
# Context switch statistics
cursor = conn.execute("""
SELECT COUNT(*), AVG(switch_cost)
FROM context_switches
WHERE timestamp > datetime('now', '-1 hour')
""")
row = cursor.fetchone()
stats['recent_switches'] = row[0] or 0
stats['avg_switch_cost'] = row[1] or 0.0
# Current focus
if self._attention_focus:
stats['current_focus'] = {
'type': self._attention_focus.focus_type,
'items': len(self._attention_focus.focus_items),
'weight': self._attention_focus.attention_weight,
'duration_seconds': (datetime.now() - self._attention_focus.created_at).total_seconds()
}
else:
stats['current_focus'] = None
return stats
except Exception as e:
logger.error(f"Error getting working memory statistics: {e}")
return {'error': str(e)}
def cleanup(self):
"""Cleanup resources"""
self._decay_running = False
if self._decay_thread.is_alive():
self._decay_thread.join(timeout=1.0)
# Export the main classes
__all__ = ['WorkingMemoryManager', 'WorkingMemoryItem', 'AttentionFocus']
|