Spaces:
Running
Running
| """Small bounded pool for periodic memory extraction; no extra service required.""" | |
| from concurrent.futures import ThreadPoolExecutor | |
| from copy import deepcopy | |
| import logging | |
| import threading | |
| from engine.memory import MAX_RECENT, update_memory | |
| logger = logging.getLogger(__name__) | |
| class MemoryScheduler: | |
| def __init__(self, workers=2): | |
| self._executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix='memory') | |
| self._capacity = workers | |
| self._pending = set() | |
| self._lock = threading.Lock() | |
| def schedule(self, session_id, character_id, conversation, client, model, persona=None): | |
| if len(conversation['recent_messages']) - conversation['summarized_count'] <= MAX_RECENT: | |
| return None | |
| key = (session_id, character_id) | |
| with self._lock: | |
| # No unbounded queue and at most one job for each conversation. | |
| # Skipped batches remain intact and are retried after a later reply. | |
| if key in self._pending or len(self._pending) >= self._capacity: | |
| return None | |
| self._pending.add(key) | |
| try: | |
| snapshot = deepcopy(conversation) | |
| future = self._executor.submit(update_memory, session_id, character_id, snapshot, client, model, deepcopy(persona)) | |
| except Exception: | |
| self._release(key) | |
| logger.warning('Could not schedule memory update; messages retained.', exc_info=True) | |
| return None | |
| future.add_done_callback(lambda _: self._release(key)) | |
| return future | |
| def _release(self, key): | |
| with self._lock: | |
| self._pending.discard(key) | |
| def shutdown(self): | |
| self._executor.shutdown(wait=True, cancel_futures=True) | |
| _scheduler = MemoryScheduler() | |
| schedule_memory_update = _scheduler.schedule | |