File size: 1,721 Bytes
add8f0b |
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 |
from __future__ import annotations
import threading
from collections import OrderedDict
from copy import deepcopy
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from gradio.blocks import Blocks
class StateHolder:
def __init__(self):
self.capacity = 10000
self.session_data = OrderedDict()
self.lock = threading.Lock()
def set_blocks(self, blocks: Blocks):
self.blocks = blocks
self.capacity = blocks.state_session_capacity
def __getitem__(self, session_id: str) -> SessionState:
if session_id not in self.session_data:
self.session_data[session_id] = SessionState(self.blocks)
self.update(session_id)
return self.session_data[session_id]
def __contains__(self, session_id: str):
return session_id in self.session_data
def update(self, session_id: str):
with self.lock:
if session_id in self.session_data:
self.session_data.move_to_end(session_id)
if len(self.session_data) > self.capacity:
self.session_data.popitem(last=False)
class SessionState:
def __init__(self, blocks: Blocks):
self.blocks = blocks
self._data = {}
def __getitem__(self, key: int) -> Any:
if key not in self._data:
block = self.blocks.blocks[key]
if getattr(block, "stateful", False):
self._data[key] = deepcopy(getattr(block, "value", None))
else:
self._data[key] = None
return self._data[key]
def __setitem__(self, key: int, value: Any):
self._data[key] = value
def __contains__(self, key: int):
return key in self._data
|