Download app/core/session/session_manager.py from bg5/zai: direct link, hf CLI and curl.
- Browser
- Download file 18.2 kB
-
https://huggingface.co/spaces/bg5/zai/resolve/main/app/core/session/session_manager.py
- Command line
-
hf download hf://spaces/bg5/zai/app/core/session/session_manager.py
-
curl -L -o session_manager.py https://huggingface.co/spaces/bg5/zai/resolve/main/app/core/session/session_manager.py
18.2 kB
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| """ไผ่ฏ็ฎก็ๅจ๏ผ่ฟ็ปญไผ่ฏ็ๆ ธๅฟๅ่ฐๆจกๅใ | |
| ่่ดฃ๏ผ | |
| 1. ้่ฟๆ็บนๅน้ ๆฃๆต่ฟ็ปญไผ่ฏ | |
| 2. ็ปดๆค chat_id / parent_id ็็ๅฝๅจๆ | |
| 3. ๅๅฐๅฎๆๆธ ็่ฟๆไผ่ฏ | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import time | |
| import uuid | |
| from dataclasses import dataclass, field | |
| from typing import Any, Dict, List, Optional | |
| from app.core.session.session_fingerprint import SessionFingerprint | |
| from app.core.session.session_store import SessionStore | |
| from app.core.config import settings | |
| from app.utils.logger import logger | |
| def _new_uuid() -> str: | |
| return str(uuid.uuid4()) | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๆฐๆฎ็ฑป | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| class SessionResult: | |
| """ไผ่ฏๆฅๆพ/ๅๅปบ็็ปๆ๏ผไพ transform_request ไฝฟ็จใ""" | |
| chat_id: str | |
| """ๅค็จๆๆฐๅปบ็ๅฏน่ฏ IDใ""" | |
| message_id: str | |
| """ๆฌๆฌก่ฏทๆฑ็ๆ็็จๆทๆถๆฏ ID๏ผcurrent_user_message_id๏ผใ""" | |
| parent_id: Optional[str] | |
| """่ฟ็ปญไผ่ฏๆถๆๅไธๆฌก็ current_user_message_id๏ผๆฐๅฏน่ฏๆถไธบ Noneใ""" | |
| is_new: bool | |
| """True = ๆฐๅปบๅฏน่ฏ๏ผFalse = ๅค็จๅทฒๆไผ่ฏใ""" | |
| bound_token: Optional[str] = None | |
| """ๅค็จไผ่ฏๆถ็ปๅฎ็ๅๅง token๏ผๆฐๅฏน่ฏๆถไธบ Noneใ""" | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๅญๅจ้ฎ่พ ๅฉ | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| _SESSION_PREFIX = "session" | |
| _FP_INDEX_PREFIX = "fp_index" | |
| def _session_key(client_fp: str, chat_id: str) -> str: | |
| return f"{_SESSION_PREFIX}:{client_fp}:{chat_id}" | |
| def _fp_index_key(client_fp: str) -> str: | |
| return f"{_FP_INDEX_PREFIX}:{client_fp}" | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # SessionManager | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| class SessionManager: | |
| """่ฟ็ปญไผ่ฏ็ฎก็ๅจใ | |
| ๅไพไฝฟ็จๆนๅผ๏ผ็ฑ UpstreamClient ๆๆ๏ผ๏ผ | |
| self._session_manager = SessionManager() | |
| """ | |
| def __init__( | |
| self, | |
| session_ttl: int = 3600, | |
| max_sessions_per_client: int = 50, | |
| cleanup_interval: int = 300, | |
| ) -> None: | |
| self._ttl = session_ttl | |
| self._max_per_client = max(1, max_sessions_per_client) | |
| self._cleanup_interval = cleanup_interval | |
| self._store = SessionStore() | |
| self._fp = SessionFingerprint() | |
| self._cleanup_task: Optional[asyncio.Task] = None | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๅ ฌๅ ฑๆฅๅฃ | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| async def find_session( | |
| self, | |
| model: str, | |
| messages: List[Dict[str, Any]], | |
| client_id: Optional[str] = None, | |
| ) -> Optional[SessionResult]: | |
| """Find matching continuous session by fingerprint. Returns None if not found. | |
| This method is intentionally read-only. It only returns the matched | |
| session and the IDs needed for the current request; the caller must | |
| commit the new current_user_message_id after the upstream request is | |
| accepted. | |
| Args: | |
| model: Model name used for index grouping. | |
| messages: Full message list (OpenAI format) for fingerprint matching. | |
| client_id: Optional explicit client ID (used for index key if provided). | |
| Returns: | |
| SessionResult on match, None otherwise. | |
| """ | |
| self._ensure_cleanup_started() | |
| index_identifier = client_id or model | |
| client_fp = self._fp.generate_client_fingerprint(index_identifier, model) | |
| match = await self._match_session(client_fp, messages) | |
| if not match: | |
| return None | |
| session = match | |
| message_id = _new_uuid() | |
| chat_id = session["chat_id"] | |
| parent_id = session.get("last_message_id") or None | |
| logger.debug( | |
| "่ฟ็ปญไผ่ฏๅค็จ: chat_id=%s... " | |
| "parent_id=%s...", | |
| chat_id[:8], | |
| parent_id[:8] if parent_id else 'None', | |
| ) | |
| return SessionResult( | |
| chat_id=chat_id, | |
| message_id=message_id, | |
| parent_id=parent_id, | |
| is_new=False, | |
| bound_token=session.get("auth_token"), | |
| ) | |
| async def create_session( | |
| self, | |
| auth_token: str, | |
| model: str, | |
| messages: List[Dict[str, Any]], | |
| chat_id: str, | |
| message_id: str, | |
| client_id: Optional[str] = None, | |
| ) -> SessionResult: | |
| """Create a new session with an externally-provided chat_id (from upstream). | |
| This is the counterpart to find_session: the caller has already obtained | |
| a chat_id from the upstream /chats/new endpoint and wants to store it | |
| for future reuse. | |
| Args: | |
| auth_token: Upstream auth token (stored for bound_token on reuse). | |
| model: Model name for index grouping. | |
| messages: Full message list (OpenAI format) for fingerprint collection. | |
| chat_id: Chat ID returned by the upstream pre-create endpoint. | |
| message_id: current_user_message_id generated for this turn. | |
| client_id: Optional explicit client ID (used for index key if provided). | |
| Returns: | |
| SessionResult with is_new=True. | |
| """ | |
| self._ensure_cleanup_started() | |
| index_identifier = client_id or model | |
| client_fp = self._fp.generate_client_fingerprint(index_identifier, model) | |
| fingerprints = self._fp.collect_fingerprints(messages) | |
| record: Dict[str, Any] = { | |
| "chat_id": chat_id, | |
| "client_fingerprint": client_fp, | |
| "fingerprints": fingerprints, | |
| "last_message_id": message_id, | |
| "model": model, | |
| "token_hash": self._fp.hash_token(auth_token), | |
| "auth_token": auth_token, | |
| "created_at": time.time(), | |
| "last_update": time.time(), | |
| } | |
| await self._save_session(client_fp, chat_id, record) | |
| await self._update_fp_index(client_fp, chat_id) | |
| logger.debug("ๆฐๅปบไผ่ฏ: chat_id=%s... model=%s", chat_id[:8], model) | |
| return SessionResult( | |
| chat_id=chat_id, | |
| message_id=message_id, | |
| parent_id=None, | |
| is_new=True, | |
| ) | |
| async def commit_session_turn( | |
| self, | |
| model: str, | |
| messages: List[Dict[str, Any]], | |
| chat_id: str, | |
| message_id: str, | |
| client_id: Optional[str] = None, | |
| ) -> None: | |
| """ๅจไธๆธธๆๅๆฅๆถ่ฏทๆฑๅ๏ผๆจ่ฟ่ฟ็ปญไผ่ฏ้พๆ้ใ | |
| Args: | |
| model: Model name for index grouping. | |
| messages: Full message list (OpenAI format) for fingerprint update. | |
| chat_id: Existing upstream chat ID. | |
| message_id: Accepted current_user_message_id for this turn. | |
| client_id: Optional explicit client ID (used for index key if provided). | |
| """ | |
| self._ensure_cleanup_started() | |
| index_identifier = client_id or model | |
| client_fp = self._fp.generate_client_fingerprint(index_identifier, model) | |
| session = await self._store.get(_session_key(client_fp, chat_id)) | |
| if not session: | |
| logger.warning( | |
| "ๆไบค่ฟ็ปญไผ่ฏๅคฑ่ดฅ: ๆชๆพๅฐ chat_id=%s ็ไผ่ฏ่ฎฐๅฝ", | |
| chat_id[:8], | |
| ) | |
| return | |
| session["fingerprints"] = self._fp.collect_fingerprints(messages) | |
| session["last_message_id"] = message_id | |
| session["last_update"] = time.time() | |
| await self._save_session(client_fp, chat_id, session) | |
| await self._touch_fp_index(client_fp, chat_id) | |
| logger.debug( | |
| "่ฟ็ปญไผ่ฏๅทฒๆไบค: chat_id=%s last_message_id=%s", | |
| chat_id[:8], | |
| message_id[:8], | |
| ) | |
| async def update_session_message_id( | |
| self, | |
| client_fp_or_chat_id: str, | |
| chat_id: str, | |
| new_message_id: str, | |
| ) -> None: | |
| """๏ผๅฏ้๏ผๆๅจๅทๆฐไผ่ฏ็ last_message_idใ | |
| ้ๅธธๆ ้่ฐ็จโโไผ่ฏ็ถๆไผๅจไธๆธธๆๅๆฅๆถ่ฏทๆฑๅๆไบคใ | |
| ๅฝ้่ฆๆพๅผไฟฎๆญฃ current_user_message_id ๆถๅฏ่ฐ็จใ | |
| """ | |
| # ๅ ๅฐ่ฏๆ chat_id ๆพ client_fp๏ผ้ๅ fp_index๏ผ | |
| # ็ฎๅ๏ผ่ฅ่ฐ็จๆนๅทฒ็ฅ client_fp ๅ็ดๆฅๆไฝ | |
| session = await self._store.get(_session_key(client_fp_or_chat_id, chat_id)) | |
| if not session: | |
| return | |
| session["last_message_id"] = new_message_id | |
| session["last_update"] = time.time() | |
| await self._store.set( | |
| _session_key(client_fp_or_chat_id, chat_id), session, ttl=self._ttl | |
| ) | |
| async def clear_session(self, client_fp: str, chat_id: str) -> None: | |
| """ๆธ ้คๅไธชไผ่ฏๅๅ ถๆ็บน็ดขๅผๆก็ฎใ""" | |
| await self._store.delete(_session_key(client_fp, chat_id)) | |
| await self._remove_from_fp_index(client_fp, chat_id) | |
| async def get_stats(self) -> Dict[str, Any]: | |
| """่ฟๅๅญๅจ็ป่ฎกไฟกๆฏ๏ผ็จไบ Admin ้ขๆฟ๏ผใ""" | |
| session_keys = await self._store.keys(f"{_SESSION_PREFIX}:") | |
| fp_keys = await self._store.keys(f"{_FP_INDEX_PREFIX}:") | |
| return { | |
| "total_sessions": len(session_keys), | |
| "total_clients": len(fp_keys), | |
| "store_size": self._store.size(), | |
| "session_ttl": self._ttl, | |
| "max_per_client": self._max_per_client, | |
| } | |
| async def close(self) -> None: | |
| """ๅๆญขๅๅฐๆธ ็ไปปๅกใ""" | |
| if self._cleanup_task and not self._cleanup_task.done(): | |
| self._cleanup_task.cancel() | |
| try: | |
| await self._cleanup_task | |
| except asyncio.CancelledError: | |
| pass | |
| async with self._store._lock: | |
| self._store._data.clear() | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๅ ้จ๏ผไผ่ฏๅน้ | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| async def _match_session( | |
| self, | |
| client_fp: str, | |
| messages: List[Dict[str, Any]], | |
| ) -> Optional[Dict[str, Any]]: | |
| """ๅจ่ฏฅๅฎขๆท็ซฏ็ไผ่ฏๅ่กจไธญๆฅๆพๅน้ ็่ฟ็ปญไผ่ฏใ""" | |
| index = await self._store.get(_fp_index_key(client_fp)) | |
| if not index: | |
| return None | |
| chat_ids: List[str] = index.get("chat_ids", []) | |
| if not chat_ids: | |
| return None | |
| # ๅพ ็กฎ่ฎคไผ่ฏ็ฒพ็กฎๅน้ ๆถๆฐๆง่ฎก็ฎ | |
| pending_fps: Optional[List[str]] = None | |
| # ไปๆๆฐ็ไผ่ฏๅผๅงๅๅบๅน้ | |
| for chat_id in reversed(chat_ids): | |
| session = await self._store.get(_session_key(client_fp, chat_id)) | |
| if not session: | |
| # ๅทฒ่ฟๆ๏ผไป็ดขๅผไธญๆธ ็ | |
| await self._remove_from_fp_index(client_fp, chat_id) | |
| continue | |
| cached_fps: List[str] = session.get("fingerprints", []) | |
| # ๅพ ็กฎ่ฎคไผ่ฏ๏ผprecreate ๆๅไฝ completions ๆช็กฎ่ฎค๏ผ๏ผ | |
| # is_continuous_session ๅฏน้ฆ่ฝฎๆถๆฏ๏ผโค2 ๆก๏ผ่ฟๅ False๏ผ | |
| # ๅ ๆญคๆน็จ็ฒพ็กฎๆ็บนๆฏๅฏนๆฅๅน้ ๅไธ่ฏทๆฑ็้่ฏใ | |
| if not session.get("last_message_id"): | |
| if pending_fps is None: | |
| pending_fps = self._fp.collect_fingerprints(messages) | |
| if pending_fps == cached_fps: | |
| return session | |
| continue | |
| if self._fp.is_continuous_session(messages, cached_fps): | |
| return session | |
| return None | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๅ ้จ๏ผๅญๅจๆไฝ | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| async def _save_session( | |
| self, client_fp: str, chat_id: str, session: Dict[str, Any] | |
| ) -> None: | |
| key = _session_key(client_fp, chat_id) | |
| await self._store.set(key, session, ttl=self._ttl) | |
| async def _update_fp_index(self, client_fp: str, chat_id: str) -> None: | |
| """ๆฐๅข chat_id ๅฐๆ็บน็ดขๅผ๏ผ่ถ ๅบไธ้ๆถ้ฉฑ้ๆๆง็ใ""" | |
| index_key = _fp_index_key(client_fp) | |
| index = await self._store.get(index_key) or { | |
| "client_fingerprint": client_fp, | |
| "chat_ids": [], | |
| "created_at": time.time(), | |
| } | |
| chat_ids: List[str] = index.get("chat_ids", []) | |
| if chat_id in chat_ids: | |
| chat_ids.remove(chat_id) | |
| chat_ids.append(chat_id) | |
| # ่ถ ๅบไธ้๏ผ้ฉฑ้ๆๆง็ไผ่ฏ | |
| while len(chat_ids) > self._max_per_client: | |
| evict_id = chat_ids.pop(0) | |
| await self._store.delete(_session_key(client_fp, evict_id)) | |
| logger.debug("้ฉฑ้่ถ ้ไผ่ฏ: %s...", evict_id[:8]) | |
| index["chat_ids"] = chat_ids | |
| index["most_recent"] = chat_id | |
| index["last_update"] = time.time() | |
| await self._store.set(index_key, index, ttl=self._ttl) | |
| async def _touch_fp_index(self, client_fp: str, chat_id: str) -> None: | |
| """ๅฐๅทฒๅน้ ็ chat_id ็งปๅฐๅ่กจๆซๅฐพ๏ผๆ่ฟไฝฟ็จ๏ผ๏ผๅนถๅทๆฐ TTLใ""" | |
| index_key = _fp_index_key(client_fp) | |
| index = await self._store.get(index_key) | |
| if not index: | |
| return | |
| chat_ids: List[str] = index.get("chat_ids", []) | |
| if chat_id in chat_ids: | |
| chat_ids.remove(chat_id) | |
| chat_ids.append(chat_id) | |
| index["chat_ids"] = chat_ids | |
| index["most_recent"] = chat_id | |
| index["last_update"] = time.time() | |
| await self._store.set(index_key, index, ttl=self._ttl) | |
| async def _remove_from_fp_index(self, client_fp: str, chat_id: str) -> None: | |
| """ไปๆ็บน็ดขๅผไธญ็งป้ค chat_id๏ผไผ่ฏๅทฒ่ฟๆๆๅทฒๅ ้ค๏ผใ""" | |
| index_key = _fp_index_key(client_fp) | |
| index = await self._store.get(index_key) | |
| if not index: | |
| return | |
| chat_ids: List[str] = index.get("chat_ids", []) | |
| if chat_id in chat_ids: | |
| chat_ids.remove(chat_id) | |
| if chat_ids: | |
| index["chat_ids"] = chat_ids | |
| index["last_update"] = time.time() | |
| await self._store.set(index_key, index, ttl=self._ttl) | |
| else: | |
| await self._store.delete(index_key) | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| # ๅ ้จ๏ผๅๅฐๆธ ็ | |
| # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def _ensure_cleanup_started(self) -> None: | |
| """ๆฐๆงๅฏๅจๅๅฐๆธ ็ๅ็จ๏ผไป ๅจไบไปถๅพช็ฏไธญ้ฆๆฌก่ฐ็จๆถๅฏๅจ๏ผใ""" | |
| if self._cleanup_task is None or self._cleanup_task.done(): | |
| try: | |
| self._cleanup_task = asyncio.get_event_loop().create_task( | |
| self._cleanup_loop() | |
| ) | |
| except RuntimeError: | |
| pass # ๆ ไบไปถๅพช็ฏ๏ผๆต่ฏ็ฏๅข๏ผๆถๅฟฝ็ฅ | |
| async def _cleanup_loop(self) -> None: | |
| """ๅฎๆๆธ ็่ฟๆไผ่ฏใ""" | |
| while True: | |
| try: | |
| await asyncio.sleep(self._cleanup_interval) | |
| removed = await self._store.cleanup_expired() | |
| if removed: | |
| logger.debug("ไผ่ฏๆธ ็๏ผๆธ ้ค %s ๆก่ฟๆ่ฎฐๅฝ", removed) | |
| except asyncio.CancelledError: | |
| break | |
| except Exception as exc: | |
| logger.warning("ไผ่ฏๆธ ็ๅผๅธธ", exc_info=settings.DEBUG_LOGGING) | |