"""State machine треда: текущая ветка, шаг внутри ветки, собранные слоты. Используется в chat_service для ведения многошаговых сценариев (Спринт 5). Слоты — произвольный JSON-словарь, конкретные ключи определяются веткой. """ import json import logging from datetime import datetime, timezone from sqlalchemy.ext.asyncio import AsyncSession from db.models import ThreadState logger = logging.getLogger(__name__) async def get_state(session: AsyncSession, thread_id: int) -> ThreadState | None: return await session.get(ThreadState, thread_id) def _parse_slots(raw: str) -> dict: if not raw: return {} try: value = json.loads(raw) except json.JSONDecodeError: logger.warning("Bad slots_json for thread_state, resetting to {}") return {} return value if isinstance(value, dict) else {} async def load_snapshot(session: AsyncSession, thread_id: int) -> dict: """Удобный снимок состояния для чтения (intent, step_code, step, slots).""" state = await get_state(session, thread_id) if state is None: return { "current_intent_code": None, "current_step": 0, "current_step_code": None, "slots": {}, } return { "current_intent_code": state.current_intent_code, "current_step": state.current_step, "current_step_code": state.current_step_code, "slots": _parse_slots(state.slots_json), } async def upsert( session: AsyncSession, thread_id: int, *, intent_code: str | None, step: int, slots: dict, step_code: str | None = None, ) -> ThreadState: """Создать или обновить состояние треда. Коммит — на совести вызывающего.""" state = await get_state(session, thread_id) now = datetime.now(timezone.utc) slots_raw = json.dumps(slots or {}, ensure_ascii=False) if state is None: state = ThreadState( thread_id=thread_id, current_intent_code=intent_code, current_step=step, current_step_code=step_code, slots_json=slots_raw, updated_at=now, ) session.add(state) else: state.current_intent_code = intent_code state.current_step = step state.current_step_code = step_code state.slots_json = slots_raw state.updated_at = now return state async def reset( session: AsyncSession, thread_id: int, *, new_intent_code: str | None, new_step_code: str | None = None, ) -> ThreadState: """Сбросить шаг и слоты треда, выставить новую ветку (при смене intent).""" return await upsert( session, thread_id, intent_code=new_intent_code, step=0, step_code=new_step_code, slots={}, )