83 lines
3.3 KiB
Python
83 lines
3.3 KiB
Python
"""NL2SQL 会话上下文的 Redis 存储与 Prompt 格式化。"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
|
|
logger = logging.getLogger("nl2sql.session_context")
|
|
|
|
|
|
class SessionContextStore:
|
|
"""按员工隔离并限制大小的短期查询上下文。"""
|
|
|
|
def __init__(self, redis, *, ttl: int = 1800, max_messages: int = 8, max_chars: int = 4000):
|
|
self.redis = redis
|
|
self.ttl = ttl
|
|
self.max_messages = max_messages
|
|
self.max_chars = max_chars
|
|
|
|
def _key(self, session_id: str) -> str:
|
|
return f"nl2sql:session:{session_id}:context"
|
|
|
|
def _owner_key(self, session_id: str) -> str:
|
|
return f"nl2sql:session:{session_id}:owner"
|
|
|
|
async def _belongs_to(self, user_id: int, session_id: str) -> bool:
|
|
owner_key = self._owner_key(session_id)
|
|
owner = await self.redis.get(owner_key)
|
|
if owner is None:
|
|
return bool(await self.redis.set(owner_key, str(user_id), nx=True, ex=self.ttl))
|
|
if isinstance(owner, bytes):
|
|
owner = owner.decode()
|
|
return str(owner) == str(user_id)
|
|
|
|
async def load(self, user_id: int, session_id: str | None) -> list[dict[str, str]]:
|
|
"""读取当前员工的上下文，异常或跨员工访问时返回空列表。"""
|
|
if not session_id:
|
|
return []
|
|
try:
|
|
if not await self._belongs_to(user_id, session_id):
|
|
return []
|
|
raw = await self.redis.get(self._key(session_id))
|
|
if not raw:
|
|
return []
|
|
if isinstance(raw, bytes):
|
|
raw = raw.decode()
|
|
value = json.loads(raw)
|
|
return value if isinstance(value, list) else []
|
|
except Exception: # noqa: BLE001 上下文故障不阻断查询
|
|
logger.warning("NL2SQL 会话上下文读取失败", exc_info=True)
|
|
return []
|
|
|
|
async def append(self, user_id: int, session_id: str | None, question: str, status: str) -> bool:
|
|
"""追加问题和状态摘要，不保存 SQL、结果行或敏感数据。"""
|
|
if not session_id:
|
|
return False
|
|
try:
|
|
if not await self._belongs_to(user_id, session_id):
|
|
return False
|
|
messages = await self.load(user_id, session_id)
|
|
messages.append({"role": "user", "content": question[:1000]})
|
|
messages.append({"role": "assistant", "content": status[:100]})
|
|
messages = messages[-self.max_messages :]
|
|
while len(json.dumps(messages, ensure_ascii=False)) > self.max_chars and messages:
|
|
messages.pop(0)
|
|
await self.redis.set(self._key(session_id), json.dumps(messages, ensure_ascii=False), ex=self.ttl)
|
|
return True
|
|
except Exception: # noqa: BLE001 上下文故障不阻断查询
|
|
logger.warning("NL2SQL 会话上下文写入失败", exc_info=True)
|
|
return False
|
|
|
|
|
|
def build_conversation_context(messages: list[dict[str, Any]], *, max_chars: int = 2000) -> str:
|
|
"""构造有界 Prompt 上下文，只保留角色和文本内容。"""
|
|
lines: list[str] = []
|
|
for message in messages:
|
|
role = str(message.get("role", ""))[:20]
|
|
content = str(message.get("content", ""))[:500]
|
|
if role and content:
|
|
lines.append(f"{role}: {content}")
|
|
return "\n".join(lines)[-max_chars:]
|