49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import threading
|
|
import time
|
|
|
|
class Session:
|
|
def __init__(self, session_id: str, system_prompt: str):
|
|
self.session_id = session_id
|
|
self.messages = [{"role": "system", "content": system_prompt}]
|
|
self.last_activity = time.monotonic()
|
|
self._timer: threading.Timer | None = None
|
|
|
|
def add_message(self, role: str, content: str, **kwargs):
|
|
msg = {"role": role, "content": content}
|
|
msg.update(kwargs)
|
|
self.messages.append(msg)
|
|
|
|
def cancel_timer(self):
|
|
if self._timer:
|
|
self._timer.cancel()
|
|
self._timer = None
|
|
|
|
def start_timer(self, timeout: float, callback, *args):
|
|
self.cancel_timer()
|
|
self._timer = threading.Timer(timeout, callback, args)
|
|
self._timer.daemon = True
|
|
self._timer.start()
|
|
|
|
class SessionManager:
|
|
def __init__(self):
|
|
self._sessions: dict[str, Session] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def get_or_create(self, session_id: str, system_prompt: str) -> Session:
|
|
with self._lock:
|
|
if session_id not in self._sessions:
|
|
self._sessions[session_id] = Session(session_id, system_prompt)
|
|
return self._sessions[session_id]
|
|
|
|
def reset(self, session_id: str):
|
|
with self._lock:
|
|
session = self._sessions.pop(session_id, None)
|
|
if session:
|
|
session.cancel_timer()
|
|
|
|
def cleanup_all(self):
|
|
with self._lock:
|
|
for session in self._sessions.values():
|
|
session.cancel_timer()
|
|
self._sessions.clear()
|