""" 进程内 NL 附属数据（会话 / 消息 / 收藏），字段与 web 端 ApiEnvelope、SessionRow、MessageRow 对齐。 仅用于本地或 demo 联调，重启后数据丢失。 """ from __future__ import annotations import asyncio import uuid from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple def utc_ts() -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") class LiteNlStore: def __init__(self) -> None: self._lock = asyncio.Lock() self._sessions: Dict[Tuple[str, str], List[Dict[str, Any]]] = {} self._messages: Dict[Tuple[str, str, str], List[Dict[str, Any]]] = {} self._msg_counters: Dict[Tuple[str, str, str], int] = {} self._fav: Dict[Tuple[str, str], Dict[str, List[Any]]] = {} @staticmethod def _visitor_key(vid: Optional[str]) -> str: return (vid or "").strip() def _scope(self, user_id: Optional[str], visitor_biz_id: Optional[str]) -> Tuple[str, str]: uid = (user_id or "anonymous").strip() or "anonymous" return uid, self._visitor_key(visitor_biz_id) async def create_session( self, user_id: Optional[str], visitor_biz_id: Optional[str], title: Optional[str], ) -> Dict[str, Any]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) sid = str(uuid.uuid4()) now = utc_ts() row: Dict[str, Any] = { "session_id": sid, "title": (title or "").strip() or None, "user_id": sk[0], "visitor_biz_id": visitor_biz_id or None, "created_at": now, "updated_at": now, } self._sessions.setdefault(sk, []).insert(0, row) self._messages[(*sk, sid)] = [] self._msg_counters[(*sk, sid)] = 0 return {"session_id": sid, "title": row["title"]} async def list_sessions( self, user_id: Optional[str], visitor_biz_id: Optional[str], limit: int, offset: int, ) -> Dict[str, Any]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) items = list(self._sessions.get(sk, [])) items.sort(key=lambda r: str(r.get("updated_at") or ""), reverse=True) sl = items[offset : offset + limit] return {"items": sl, "limit": limit, "offset": offset} def _bump_session(self, sk: Tuple[str, str], session_id: str, user_first_line: str) -> None: now = utc_ts() for r in self._sessions.get(sk, []): if r.get("session_id") != session_id: continue r["updated_at"] = now if not (r.get("title") or "").strip() and user_first_line.strip(): r["title"] = user_first_line.strip()[:80] break async def append_exchange( self, user_id: Optional[str], visitor_biz_id: Optional[str], session_id: str, user_text: str, assistant_content_json: str, ) -> None: async with self._lock: sk = self._scope(user_id, visitor_biz_id) sk3 = (*sk, session_id) if sk3 not in self._messages: self._messages[sk3] = [] self._msg_counters[sk3] = 0 now = utc_ts() def next_id() -> int: n = self._msg_counters.get(sk3, 0) + 1 self._msg_counters[sk3] = n return n self._messages[sk3].append( { "id": next_id(), "role": "user", "content": user_text, "created_at": now, "llm_total_tokens": None, } ) self._messages[sk3].append( { "id": next_id(), "role": "assistant", "content": assistant_content_json, "created_at": now, "llm_total_tokens": None, } ) self._bump_session(sk, session_id, user_text) async def get_messages( self, user_id: Optional[str], visitor_biz_id: Optional[str], session_id: str, limit: int, offset: int, ) -> Optional[Dict[str, Any]]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) sk3 = (*sk, session_id) if sk3 not in self._messages: found = any( r.get("session_id") == session_id for r in self._sessions.get(sk, []) ) if not found: return None self._messages[sk3] = [] self._msg_counters[sk3] = 0 rows = list(self._messages[sk3]) rows = rows[offset : offset + limit] return {"session_id": session_id, "items": rows, "limit": limit, "offset": offset} async def update_session_title( self, user_id: Optional[str], visitor_biz_id: Optional[str], session_id: str, title: Optional[str], ) -> Optional[Dict[str, Any]]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) for r in self._sessions.get(sk, []): if r.get("session_id") == session_id: r["title"] = title r["updated_at"] = utc_ts() return {"session_id": session_id, "title": r.get("title")} return None async def delete_session( self, user_id: Optional[str], visitor_biz_id: Optional[str], session_id: str, ) -> Optional[Dict[str, Any]]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) lst = self._sessions.get(sk, []) sk3 = (*sk, session_id) n = len(self._messages.get(sk3, [])) new_lst = [x for x in lst if x.get("session_id") != session_id] if len(new_lst) == len(lst): return None self._sessions[sk] = new_lst self._messages.pop(sk3, None) self._msg_counters.pop(sk3, None) return {"session_id": session_id, "deleted_messages": n} async def patch_message( self, user_id: Optional[str], visitor_biz_id: Optional[str], session_id: str, message_id: int, content: str, ) -> Optional[Dict[str, Any]]: async with self._lock: sk3 = (*self._scope(user_id, visitor_biz_id), session_id) for m in self._messages.get(sk3, []): if m.get("id") == message_id: m["content"] = content m["created_at"] = utc_ts() return { "id": message_id, "session_id": session_id, "role": m.get("role"), "content": content, "created_at": m.get("created_at"), "llm_total_tokens": m.get("llm_total_tokens"), } return None async def get_favorites_grouped( self, user_id: Optional[str], visitor_biz_id: Optional[str] ) -> Dict[str, List[Any]]: async with self._lock: sk = self._scope(user_id, visitor_biz_id) g = self._fav.get(sk) or {"sql": [], "function": [], "report": []} return {"sql": list(g["sql"]), "function": list(g["function"]), "report": list(g["report"])} async def add_favorite( self, user_id: Optional[str], visitor_biz_id: Optional[str], body: Dict[str, Any] ) -> Any: async with self._lock: sk = self._scope(user_id, visitor_biz_id) g = self._fav.setdefault(sk, {"sql": [], "function": [], "report": []}) fav_type = str(body.get("fav_type") or "sql") fid = f"{uuid.uuid4().hex[:12]}" if fav_type == "sql": row = { "id": fid, "name": str(body.get("name") or ""), "desc": str(body.get("desc") or ""), "sql": str(body.get("sql") or ""), } if body.get("sql_explain"): row["sql_explain"] = str(body["sql_explain"]) g["sql"].insert(0, row) return row if fav_type == "function": row = { "id": fid, "name": str(body.get("name") or ""), "desc": str(body.get("desc") or ""), "path": str(body.get("path") or ""), } g["function"].insert(0, row) return row row = { "id": fid, "name": str(body.get("name") or ""), "desc": str(body.get("desc") or ""), "reportPath": str(body.get("reportPath") or ""), "params": str(body.get("params") or ""), } g["report"].insert(0, row) return row async def patch_favorite( self, user_id: Optional[str], visitor_biz_id: Optional[str], fav_id: str, patch: Dict[str, Any], ) -> bool: async with self._lock: sk = self._scope(user_id, visitor_biz_id) g = self._fav.get(sk) if not g: return False for bucket in ("sql", "function", "report"): lst = g.get(bucket, []) for i, it in enumerate(lst): if str(it.get("id")) == fav_id: lst[i] = {**it, **patch} return True return False async def delete_favorite( self, user_id: Optional[str], visitor_biz_id: Optional[str], fav_id: str ) -> bool: async with self._lock: sk = self._scope(user_id, visitor_biz_id) g = self._fav.get(sk) if not g: return False removed = False for bucket in ("sql", "function", "report"): before = len(g.get(bucket, [])) g[bucket] = [x for x in g.get(bucket, []) if str(x.get("id")) != fav_id] if len(g[bucket]) < before: removed = True return removed lite_nl_store = LiteNlStore()