"""Private knowledge-base Q&A sessions with shared Markdown publications. This persistence unit deliberately contains no Gateway/application imports. A session and its LangGraph thread belong to exactly one user. Once published, only the saved Markdown snapshot is shared with readers of the bound knowledge base; the private draft and thread identifier never become publication data. """ from __future__ import annotations import asyncio from collections.abc import AsyncIterator from contextlib import asynccontextmanager from copy import deepcopy from datetime import UTC, datetime from typing import Any from uuid import uuid4 from sqlalchemy import String, delete, select, update from sqlalchemy.exc import OperationalError, ProgrammingError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import Mapped, mapped_column from deerflow.persistence.base import Base from deerflow.persistence.types import BeijingDateTime, PortableLongText class KnowledgeResearchSessionRow(Base): __tablename__ = "knowledge_research_sessions" id: Mapped[str] = mapped_column(String(36), primary_key=True) # A separate opaque id is the only identity exposed on the shared bookshelf. publication_id: Mapped[str] = mapped_column(String(36), unique=True, index=True) thread_id: Mapped[str] = mapped_column(String(128), unique=True, index=True) knowledge_base_id: Mapped[str] = mapped_column(String(128), index=True) user_id: Mapped[str] = mapped_column(String(255), index=True) title: Mapped[str] = mapped_column(String(255)) draft_markdown: Mapped[str] = mapped_column(PortableLongText, default="") wiki_slug: Mapped[str | None] = mapped_column(String(512), nullable=True) wiki_page_id: Mapped[str | None] = mapped_column(String(128), nullable=True) wiki_folder_id: Mapped[str | None] = mapped_column(String(128), nullable=True) saved_title: Mapped[str | None] = mapped_column(String(255), nullable=True) saved_markdown: Mapped[str | None] = mapped_column(PortableLongText, nullable=True) saved_hash: Mapped[str | None] = mapped_column(String(64), nullable=True) saved_at: Mapped[datetime | None] = mapped_column(BeijingDateTime(), nullable=True) created_at: Mapped[datetime] = mapped_column( BeijingDateTime(), default=lambda: datetime.now(UTC), nullable=False, ) updated_at: Mapped[datetime] = mapped_column( BeijingDateTime(), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC), nullable=False, ) def _is_missing_table(exc: Exception) -> bool: message = str(exc).lower() return ( "no such table" in message or "doesn't exist" in message or "does not exist" in message or "unknown table" in message ) def _iso(value: Any) -> Any: return value.isoformat() if isinstance(value, datetime) else value def _row_dict(row: KnowledgeResearchSessionRow | dict[str, Any]) -> dict[str, Any]: if isinstance(row, dict): return deepcopy(row) return { "id": row.id, "publication_id": row.publication_id, "thread_id": row.thread_id, "knowledge_base_id": row.knowledge_base_id, "user_id": row.user_id, "title": row.title, "draft_markdown": row.draft_markdown or "", "wiki_slug": row.wiki_slug, "wiki_page_id": row.wiki_page_id, "wiki_folder_id": row.wiki_folder_id, "saved_title": row.saved_title, "saved_markdown": row.saved_markdown, "saved_hash": row.saved_hash, "saved_at": _iso(row.saved_at), "created_at": _iso(row.created_at), "updated_at": _iso(row.updated_at), } class _PublicationLease: """One transaction-scoped mutable publication row.""" def __init__(self, row: KnowledgeResearchSessionRow | dict[str, Any]) -> None: self._row = row def snapshot(self) -> dict[str, Any]: return _row_dict(self._row) def save( self, *, title: str, markdown: str, saved_hash: str, wiki_slug: str, wiki_page_id: str | None, wiki_folder_id: str | None, ) -> dict[str, Any]: now = datetime.now(UTC) values = { "wiki_slug": wiki_slug, "wiki_page_id": wiki_page_id, "wiki_folder_id": wiki_folder_id, "saved_title": title, "saved_markdown": markdown, "saved_hash": saved_hash, "saved_at": now, "updated_at": now, } if isinstance(self._row, dict): self._row.update( { **values, "saved_at": now.isoformat(), "updated_at": now.isoformat(), } ) else: for key, value in values.items(): setattr(self._row, key, value) return self.snapshot() class KnowledgeResearchStore: """SQL store with a behaviorally equivalent memory fallback.""" def __init__( self, session_factory: async_sessionmaker[AsyncSession] | None, ) -> None: self._sf = session_factory self._rows: dict[str, dict[str, Any]] = {} self._publication_locks: dict[str, asyncio.Lock] = {} @asynccontextmanager async def locked_publication( self, session_id: str, user_id: str, ) -> AsyncIterator[_PublicationLease | None]: """Serialize version-check + remote write + local publication commit. SQL deployments hold a row lock for the context lifetime, so separate application workers cannot both pass the same expected-hash check. SQLite has no ``FOR UPDATE`` support, so a no-op update first acquires its database write lock. The memory fallback uses one asyncio lock per research session and commits a copy only on successful exit. """ if self._sf is None: lock = self._publication_locks.setdefault(session_id, asyncio.Lock()) async with lock: original = self._rows.get(session_id) if original is None or original["user_id"] != user_id: yield None return working = deepcopy(original) lease = _PublicationLease(working) yield lease self._rows[session_id] = working return async with self._sf() as session: async with session.begin(): bind = session.get_bind() if bind.dialect.name == "sqlite": await session.execute( update(KnowledgeResearchSessionRow) .where( KnowledgeResearchSessionRow.id == session_id, KnowledgeResearchSessionRow.user_id == user_id, ) .values( updated_at=KnowledgeResearchSessionRow.updated_at ) ) row = ( await session.execute( select(KnowledgeResearchSessionRow) .where( KnowledgeResearchSessionRow.id == session_id, KnowledgeResearchSessionRow.user_id == user_id, ) .with_for_update() ) ).scalar_one_or_none() yield _PublicationLease(row) if row is not None else None async def create( self, *, knowledge_base_id: str, user_id: str, title: str, session_id: str | None = None, publication_id: str | None = None, thread_id: str | None = None, ) -> dict[str, Any]: now = datetime.now(UTC) row = KnowledgeResearchSessionRow( id=session_id or str(uuid4()), publication_id=publication_id or str(uuid4()), thread_id=thread_id or str(uuid4()), knowledge_base_id=knowledge_base_id, user_id=user_id, title=title, draft_markdown="", created_at=now, updated_at=now, ) if self._sf is None: data = _row_dict(row) self._rows[row.id] = data return deepcopy(data) async with self._sf() as session: session.add(row) await session.commit() return _row_dict(row) async def get(self, session_id: str) -> dict[str, Any] | None: if self._sf is None: row = self._rows.get(session_id) return deepcopy(row) if row else None async with self._sf() as session: row = await session.get(KnowledgeResearchSessionRow, session_id) return _row_dict(row) if row else None async def get_owned(self, session_id: str, user_id: str) -> dict[str, Any] | None: row = await self.get(session_id) return row if row and row["user_id"] == user_id else None async def get_by_thread(self, thread_id: str) -> dict[str, Any] | None: if self._sf is None: for row in self._rows.values(): if row["thread_id"] == thread_id: return deepcopy(row) return None async with self._sf() as session: try: row = ( await session.execute( select(KnowledgeResearchSessionRow).where( KnowledgeResearchSessionRow.thread_id == thread_id ) ) ).scalar_one_or_none() except (OperationalError, ProgrammingError) as exc: # Ordinary Q&A hits this lookup on every run. A missing table # (older DB before create_all registered the model) must not # abort /runs/stream; treat it as "not a research thread". if _is_missing_table(exc): return None raise return _row_dict(row) if row else None async def get_publication(self, publication_id: str) -> dict[str, Any] | None: if self._sf is None: for row in self._rows.values(): if row["publication_id"] == publication_id and row.get("saved_hash"): return deepcopy(row) return None async with self._sf() as session: row = ( await session.execute( select(KnowledgeResearchSessionRow).where( KnowledgeResearchSessionRow.publication_id == publication_id, KnowledgeResearchSessionRow.saved_hash.is_not(None), ) ) ).scalar_one_or_none() return _row_dict(row) if row else None async def list_owned( self, knowledge_base_id: str, user_id: str, ) -> list[dict[str, Any]]: if self._sf is None: rows = [ deepcopy(row) for row in self._rows.values() if row["knowledge_base_id"] == knowledge_base_id and row["user_id"] == user_id ] return sorted(rows, key=lambda item: str(item["updated_at"]), reverse=True) async with self._sf() as session: rows = ( await session.execute( select(KnowledgeResearchSessionRow) .where( KnowledgeResearchSessionRow.knowledge_base_id == knowledge_base_id, KnowledgeResearchSessionRow.user_id == user_id, ) .order_by(KnowledgeResearchSessionRow.updated_at.desc()) ) ).scalars() return [_row_dict(row) for row in rows] async def list_published( self, knowledge_base_id: str, ) -> list[dict[str, Any]]: if self._sf is None: rows = [ deepcopy(row) for row in self._rows.values() if row["knowledge_base_id"] == knowledge_base_id and row.get("saved_hash") ] return sorted( rows, key=lambda item: str(item.get("saved_at") or item["updated_at"]), reverse=True, ) async with self._sf() as session: rows = ( await session.execute( select(KnowledgeResearchSessionRow) .where( KnowledgeResearchSessionRow.knowledge_base_id == knowledge_base_id, KnowledgeResearchSessionRow.saved_hash.is_not(None), ) .order_by(KnowledgeResearchSessionRow.saved_at.desc()) ) ).scalars() return [_row_dict(row) for row in rows] async def update_owned( self, session_id: str, user_id: str, *, title: str | None = None, draft_markdown: str | None = None, ) -> dict[str, Any] | None: now = datetime.now(UTC) if self._sf is None: lock = self._publication_locks.setdefault(session_id, asyncio.Lock()) async with lock: row = self._rows.get(session_id) if row is None or row["user_id"] != user_id: return None if title is not None: row["title"] = title if draft_markdown is not None: row["draft_markdown"] = draft_markdown row["updated_at"] = now.isoformat() return deepcopy(row) async with self._sf() as session: row = ( await session.execute( select(KnowledgeResearchSessionRow) .where( KnowledgeResearchSessionRow.id == session_id, KnowledgeResearchSessionRow.user_id == user_id, ) .with_for_update() ) ).scalar_one_or_none() if row is None: return None if title is not None: row.title = title if draft_markdown is not None: row.draft_markdown = draft_markdown row.updated_at = now await session.commit() return _row_dict(row) async def save_publication( self, session_id: str, user_id: str, *, title: str, markdown: str, saved_hash: str, wiki_slug: str, wiki_page_id: str | None, wiki_folder_id: str | None, ) -> dict[str, Any] | None: async with self.locked_publication(session_id, user_id) as publication: if publication is None: return None return publication.save( title=title, markdown=markdown, saved_hash=saved_hash, wiki_slug=wiki_slug, wiki_page_id=wiki_page_id, wiki_folder_id=wiki_folder_id, ) async def delete(self, session_id: str) -> None: if self._sf is None: lock = self._publication_locks.setdefault(session_id, asyncio.Lock()) async with lock: self._rows.pop(session_id, None) return async with self._sf() as session: await session.execute( delete(KnowledgeResearchSessionRow).where( KnowledgeResearchSessionRow.id == session_id ) ) await session.commit() def make_knowledge_research_store(session_factory): return KnowledgeResearchStore(session_factory) __all__ = [ "KnowledgeResearchSessionRow", "KnowledgeResearchStore", "make_knowledge_research_store", ]