438 lines
16 KiB
Python
438 lines
16 KiB
Python
"""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",
|
|
]
|