"""SQLAlchemy implementation of the local Wiki index store.""" from __future__ import annotations from collections.abc import AsyncIterator from datetime import UTC, datetime, timedelta from typing import Any from uuid import uuid4 from sqlalchemy import and_, case, delete, func, insert, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import aliased from deerflow.persistence.llmwiki_index.base import LlmWikiIndexStore from deerflow.persistence.llmwiki_index.model import ( LlmWikiWikiEdgeRow, LlmWikiWikiPageRow, LlmWikiWikiSyncStateRow, LlmWikiWikiVectorRow, ) def _now() -> datetime: return datetime.now(UTC) def _dict(row) -> dict[str, Any]: value = row.to_dict() for key, item in tuple(value.items()): if isinstance(item, datetime): value[key] = item.isoformat() return value class SqlLlmWikiIndexStore(LlmWikiIndexStore): def __init__(self, session_factory: async_sessionmaker[AsyncSession], *, write_batch_size: int = 50) -> None: self._sf = session_factory self.write_batch_size = max(1, min(int(write_batch_size), 500)) async def _ensure_state(self, session: AsyncSession, mapping_id: str) -> LlmWikiWikiSyncStateRow: row = await session.get(LlmWikiWikiSyncStateRow, mapping_id) if row is None: row = LlmWikiWikiSyncStateRow(knowledge_base_mapping_id=mapping_id, state="idle", index_revision=0, updated_at=_now()) session.add(row) await session.flush() return row async def _fence_current_sync_lease( self, session: AsyncSession, mapping_id: str, *, owner: str, sync_id: str, ) -> bool: """Fence destructive snapshot writes behind the current live lease.""" # Use a conditional UPDATE rather than SELECT ... FOR UPDATE. SQLite # ignores FOR UPDATE, while this statement obtains its write lock on # every supported backend before the page/edge mutation starts. stmt = ( update(LlmWikiWikiSyncStateRow) .where( LlmWikiWikiSyncStateRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiSyncStateRow.lease_owner == owner, LlmWikiWikiSyncStateRow.active_sync_id == sync_id, LlmWikiWikiSyncStateRow.lease_expires_at.is_not(None), LlmWikiWikiSyncStateRow.lease_expires_at >= _now(), ) .values(updated_at=_now()) ) return bool((await session.execute(stmt)).rowcount) async def get_page_by_slug_hash(self, mapping_id: str, slug_hash: str) -> dict[str, Any] | None: count = select(func.count(LlmWikiWikiVectorRow.id)).where(LlmWikiWikiVectorRow.page_id == LlmWikiWikiPageRow.id).correlate(LlmWikiWikiPageRow).scalar_subquery() stmt = select(LlmWikiWikiPageRow, count).where(LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.slug_hash == slug_hash) async with self._sf() as session: result = (await session.execute(stmt)).one_or_none() return {**_dict(result[0]), "vector_count": result[1]} if result else None async def get_page(self, mapping_id: str, slug: str) -> dict[str, Any] | None: stmt = select(LlmWikiWikiPageRow).where(LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.slug == slug) async with self._sf() as session: row = (await session.execute(stmt)).scalar_one_or_none() return _dict(row) if row else None async def search_pages_by_keywords( self, mapping_ids: list[str], tokens: list[str], *, limit: int = 100, include_drafts: bool = False, ) -> list[dict[str, Any]]: normalized_ids = list(dict.fromkeys(str(item).strip() for item in mapping_ids if str(item).strip())) normalized_tokens = list(dict.fromkeys(str(item).strip().lower() for item in tokens if str(item).strip())) if not normalized_ids or not normalized_tokens: return [] conditions = [] for token in normalized_tokens[:32]: needle = f"%{token}%" conditions.extend( ( func.lower(LlmWikiWikiPageRow.title).like(needle), func.lower(LlmWikiWikiPageRow.summary).like(needle), func.lower(LlmWikiWikiPageRow.content_md).like(needle), ) ) statuses = ["published", "draft"] if include_drafts else ["published"] stmt = ( select(LlmWikiWikiPageRow) .where( LlmWikiWikiPageRow.knowledge_base_mapping_id.in_(normalized_ids), LlmWikiWikiPageRow.is_remote_deleted.is_(False), LlmWikiWikiPageRow.status.in_(statuses), func.length(func.trim(LlmWikiWikiPageRow.content_md)) > 0, or_(*conditions), ) .order_by(LlmWikiWikiPageRow.updated_at.desc()) .limit(max(1, min(int(limit), 2000))) ) async with self._sf() as session: return [_dict(row) for row in (await session.execute(stmt)).scalars()] async def list_pages( self, mapping_ids: list[str], *, limit: int = 500, include_drafts: bool = False, ) -> list[dict[str, Any]]: normalized_ids = list( dict.fromkeys(str(item).strip() for item in mapping_ids if str(item).strip()) ) if not normalized_ids: return [] statuses = ["published", "draft"] if include_drafts else ["published"] stmt = ( select(LlmWikiWikiPageRow) .where( LlmWikiWikiPageRow.knowledge_base_mapping_id.in_(normalized_ids), LlmWikiWikiPageRow.is_remote_deleted.is_(False), LlmWikiWikiPageRow.status.in_(statuses), ) .order_by( LlmWikiWikiPageRow.remote_updated_at.desc(), LlmWikiWikiPageRow.updated_at.desc(), ) .limit(max(1, min(int(limit), 2000))) ) async with self._sf() as session: return [_dict(row) for row in (await session.execute(stmt)).scalars()] async def replace_graph_edges( self, mapping_id: str, edges: list[dict[str, Any]], *, owner: str, sync_id: str, ) -> dict[str, Any] | None: """Atomically replace an authorized mapping's resolved edge snapshot.""" async with self._sf() as session: if not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return None page_stmt = select(LlmWikiWikiPageRow).where( LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.is_remote_deleted.is_(False), LlmWikiWikiPageRow.status.in_(["published", "draft"]), ) pages = list((await session.execute(page_stmt)).scalars()) pages_by_slug = {page.slug: page for page in pages} canonical = { ( str(edge.get("source_slug") or "").strip(), str(edge.get("target_slug") or "").strip(), str(edge.get("edge_type") or "wiki_link").strip()[:32] or "wiki_link", ) for edge in edges if isinstance(edge, dict) } canonical = {edge for edge in canonical if edge[0] in pages_by_slug and edge[1] in pages_by_slug and edge[0] != edge[1]} existing_stmt = select( LlmWikiWikiEdgeRow.source_page_id, LlmWikiWikiEdgeRow.target_page_id, LlmWikiWikiEdgeRow.edge_type, ).where(LlmWikiWikiEdgeRow.knowledge_base_mapping_id == mapping_id) existing = {(str(source_id), str(target_id), str(edge_type)) for source_id, target_id, edge_type in (await session.execute(existing_stmt)).all()} replacement_ids = { ( pages_by_slug[source].id, pages_by_slug[target].id, edge_type, ) for source, target, edge_type in canonical } changed = existing != replacement_ids if changed: await session.execute(delete(LlmWikiWikiEdgeRow).where(LlmWikiWikiEdgeRow.knowledge_base_mapping_id == mapping_id)) now = _now() edge_rows = [ { "id": str(uuid4()), "knowledge_base_mapping_id": mapping_id, "source_page_id": pages_by_slug[source].id, "target_page_id": pages_by_slug[target].id, "edge_type": edge_type, "created_at": now, } for source, target, edge_type in sorted(canonical) ] for start in range(0, len(edge_rows), self.write_batch_size): await session.execute( insert(LlmWikiWikiEdgeRow), edge_rows[start : start + self.write_batch_size], ) state = await self._ensure_state(session, mapping_id) state.index_revision = int(state.index_revision or 0) + 1 state.updated_at = now await session.commit() return {"edge_count": len(canonical), "changed": changed} async def get_adjacent_pages( self, mapping_ids: list[str], seed_slugs: dict[str, list[str]] | list[str], *, direction: str = "both", limit: int = 100, include_drafts: bool = False, ) -> list[dict[str, Any]]: allowed = list(dict.fromkeys(str(item).strip() for item in mapping_ids if str(item).strip())) if not allowed: return [] if direction not in {"incoming", "outgoing", "both"}: raise ValueError("direction must be incoming, outgoing, or both") if isinstance(seed_slugs, dict): seeds = {mapping_id: list(dict.fromkeys(str(slug).strip() for slug in seed_slugs.get(mapping_id, []) if str(slug).strip())) for mapping_id in allowed} else: shared = list(dict.fromkeys(str(slug).strip() for slug in seed_slugs if str(slug).strip())) seeds = {mapping_id: shared for mapping_id in allowed} seeds = {mapping_id: slugs for mapping_id, slugs in seeds.items() if slugs} if not seeds: return [] source_page = aliased(LlmWikiWikiPageRow) target_page = aliased(LlmWikiWikiPageRow) statuses = ["published", "draft"] if include_drafts else ["published"] cap = max(1, min(int(limit), 1000)) async def _query(traversal_direction: str) -> list[dict[str, Any]]: seed_page = source_page if traversal_direction == "outgoing" else target_page neighbor_page = target_page if traversal_direction == "outgoing" else source_page seed_conditions = or_( *( and_( LlmWikiWikiEdgeRow.knowledge_base_mapping_id == mapping_id, seed_page.slug.in_(slugs), ) for mapping_id, slugs in seeds.items() ) ) stmt = ( select(LlmWikiWikiEdgeRow, source_page, target_page) .join(source_page, source_page.id == LlmWikiWikiEdgeRow.source_page_id) .join(target_page, target_page.id == LlmWikiWikiEdgeRow.target_page_id) .where( LlmWikiWikiEdgeRow.knowledge_base_mapping_id.in_(allowed), source_page.knowledge_base_mapping_id == LlmWikiWikiEdgeRow.knowledge_base_mapping_id, target_page.knowledge_base_mapping_id == LlmWikiWikiEdgeRow.knowledge_base_mapping_id, source_page.is_remote_deleted.is_(False), target_page.is_remote_deleted.is_(False), source_page.status.in_(statuses), target_page.status.in_(statuses), seed_conditions, ) .order_by( LlmWikiWikiEdgeRow.knowledge_base_mapping_id, seed_page.slug, neighbor_page.slug, ) .limit(cap) ) async with self._sf() as session: triples = list((await session.execute(stmt)).all()) rows: list[dict[str, Any]] = [] for edge, source, target in triples: seed = source if traversal_direction == "outgoing" else target neighbor = target if traversal_direction == "outgoing" else source rows.append( { "knowledge_base_mapping_id": edge.knowledge_base_mapping_id, "seed_slug": seed.slug, "neighbor_slug": neighbor.slug, "direction": traversal_direction, "edge_type": edge.edge_type, "source_slug": source.slug, "target_slug": target.slug, "page": _dict(neighbor), } ) return rows result: list[dict[str, Any]] = [] if direction in {"outgoing", "both"}: result.extend(await _query("outgoing")) if direction in {"incoming", "both"}: result.extend(await _query("incoming")) result.sort( key=lambda row: ( row["knowledge_base_mapping_id"], row["seed_slug"], row["neighbor_slug"], row["direction"], ) ) deduplicated: list[dict[str, Any]] = [] seen: set[tuple[str, str, str, str]] = set() for row in result: key = ( row["knowledge_base_mapping_id"], row["seed_slug"], row["neighbor_slug"], row["direction"], ) if key in seen: continue seen.add(key) deduplicated.append(row) if len(deduplicated) >= cap: break return deduplicated async def touch_page( self, page_id: str, sync_id: str, *, mapping_id: str | None = None, owner: str | None = None, ) -> bool: async with self._sf() as session: if owner is not None: if not mapping_id or not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return False conditions = [LlmWikiWikiPageRow.id == page_id] if mapping_id is not None: conditions.append(LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id) result = await session.execute( update(LlmWikiWikiPageRow) .where(*conditions) .values(last_seen_sync_id=sync_id, last_error=None, updated_at=_now()) ) await session.commit() return bool(result.rowcount) async def refresh_page_metadata( self, *, page: dict[str, Any], sync_id: str, owner: str, ) -> dict[str, Any] | None: mapping_id = str(page["knowledge_base_mapping_id"]) async with self._sf() as session: if not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return None stmt = ( select(LlmWikiWikiPageRow) .where( LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.slug_hash == page["slug_hash"], ) .with_for_update() ) row = (await session.execute(stmt)).scalar_one_or_none() if row is None: return None if row.slug != page["slug"]: raise ValueError("Wiki slug hash collision detected") durable_changed = bool(row.is_remote_deleted or row.last_error) or any( getattr(row, key) != value for key, value in page.items() ) for key, value in page.items(): setattr(row, key, value) now = _now() row.last_seen_sync_id = sync_id row.last_error = None row.is_remote_deleted = False row.updated_at = now if durable_changed: state = await self._ensure_state(session, mapping_id) state.index_revision = int(state.index_revision or 0) + 1 state.updated_at = now await session.commit() return {"changed": durable_changed, "page": _dict(row)} async def replace_page_vectors( self, *, page: dict[str, Any], vectors: list[dict[str, Any]], fingerprint: str, sync_id: str, owner: str | None = None, ) -> dict[str, Any] | None: mapping_id = str(page["knowledge_base_mapping_id"]) async with self._sf() as session: if owner is not None and not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return None stmt = select(LlmWikiWikiPageRow).where(LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.slug_hash == page["slug_hash"]).with_for_update() row = (await session.execute(stmt)).scalar_one_or_none() if row is not None and row.slug != page["slug"]: raise ValueError("Wiki slug hash collision detected") now = _now() if row is None: row = LlmWikiWikiPageRow(id=str(uuid4()), created_at=now, **page) session.add(row) else: for key, value in page.items(): setattr(row, key, value) row.index_status = "ready" row.indexed_fingerprint = fingerprint row.indexed_at = now row.last_error = None row.is_remote_deleted = False row.last_seen_sync_id = sync_id row.updated_at = now await session.flush() await session.execute(delete(LlmWikiWikiVectorRow).where(LlmWikiWikiVectorRow.page_id == row.id)) vector_rows = [ { "id": str(uuid4()), "page_id": row.id, "knowledge_base_mapping_id": mapping_id, "section_index": int(item["section_index"]), "heading": item.get("heading"), "section_content": str(item["section_content"]), "content_hash": str(item["content_hash"]), "embedding_fingerprint": fingerprint, "embedding_dimensions": int(item["embedding_dimensions"]), "vector_blob": bytes(item["vector_blob"]), "created_at": now, } for item in vectors ] # Keep the page replacement atomic, but cap each SQL statement so # large Wiki pages cannot flood the database with one huge insert. for start in range(0, len(vector_rows), self.write_batch_size): await session.execute( insert(LlmWikiWikiVectorRow), vector_rows[start : start + self.write_batch_size], ) state = await self._ensure_state(session, mapping_id) state.index_revision = int(state.index_revision or 0) + 1 state.embedding_fingerprint = fingerprint state.updated_at = now await session.commit() return _dict(row) async def record_page_failure( self, *, page: dict[str, Any], sync_id: str, error: str, owner: str | None = None, ) -> bool: mapping_id = str(page["knowledge_base_mapping_id"]) safe_error = error[:4000] async with self._sf() as session: if owner is not None and not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return False stmt = select(LlmWikiWikiPageRow).where(LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.slug_hash == page["slug_hash"]) row = (await session.execute(stmt)).scalar_one_or_none() now = _now() if row is None: row = LlmWikiWikiPageRow(id=str(uuid4()), created_at=now, **page) session.add(row) row.index_status = "failed" else: vector_count = int((await session.execute(select(func.count()).select_from(LlmWikiWikiVectorRow).where(LlmWikiWikiVectorRow.page_id == row.id))).scalar_one()) # Preserve the last atomically completed generation while the # newer remote version remains pending after an embedding error. row.index_status = "ready" if vector_count else "failed" row.last_seen_sync_id = sync_id row.last_error = safe_error row.updated_at = now await self._ensure_state(session, mapping_id) await session.commit() return True async def mark_missing_pages_deleted( self, mapping_id: str, sync_id: str, *, owner: str, ) -> int | None: async with self._sf() as session: if not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return None stmt = ( select(LlmWikiWikiPageRow) .where( LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.is_remote_deleted.is_(False), or_(LlmWikiWikiPageRow.last_seen_sync_id.is_(None), LlmWikiWikiPageRow.last_seen_sync_id != sync_id), ) .with_for_update() ) rows = list((await session.execute(stmt)).scalars()) if not rows: return 0 now = _now() page_ids = [row.id for row in rows] await session.execute(delete(LlmWikiWikiVectorRow).where(LlmWikiWikiVectorRow.page_id.in_(page_ids))) for row in rows: row.is_remote_deleted = True row.index_status = "deleted" row.updated_at = now state = await self._ensure_state(session, mapping_id) state.index_revision = int(state.index_revision or 0) + 1 state.updated_at = now await session.commit() return len(rows) async def load_vector_snapshot(self, mapping_id: str, fingerprint: str) -> dict[str, Any]: stmt = ( select(LlmWikiWikiVectorRow, LlmWikiWikiPageRow) .join(LlmWikiWikiPageRow, LlmWikiWikiPageRow.id == LlmWikiWikiVectorRow.page_id) .where( LlmWikiWikiVectorRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiVectorRow.embedding_fingerprint == fingerprint, LlmWikiWikiPageRow.index_status == "ready", LlmWikiWikiPageRow.is_remote_deleted.is_(False), LlmWikiWikiPageRow.status.in_(["published", "draft"]), ) .order_by(LlmWikiWikiVectorRow.page_id, LlmWikiWikiVectorRow.section_index) ) async with self._sf() as session: state = await session.get(LlmWikiWikiSyncStateRow, mapping_id) pairs = list((await session.execute(stmt)).all()) rows = [{**_dict(vector), "page": _dict(page)} for vector, page in pairs] return {"revision": int(state.index_revision if state else 0), "fingerprint": fingerprint, "rows": rows} async def iter_vector_snapshot_rows( self, mapping_id: str, fingerprint: str, *, batch_size: int = 1000, ) -> AsyncIterator[dict[str, Any]]: batch = max(1, min(int(batch_size), 5000)) offset = 0 while True: stmt = ( select(LlmWikiWikiVectorRow, LlmWikiWikiPageRow) .join(LlmWikiWikiPageRow, LlmWikiWikiPageRow.id == LlmWikiWikiVectorRow.page_id) .where( LlmWikiWikiVectorRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiVectorRow.embedding_fingerprint == fingerprint, LlmWikiWikiPageRow.index_status == "ready", LlmWikiWikiPageRow.is_remote_deleted.is_(False), LlmWikiWikiPageRow.status.in_(["published", "draft"]), ) .order_by(LlmWikiWikiVectorRow.page_id, LlmWikiWikiVectorRow.section_index) .limit(batch) .offset(offset) ) async with self._sf() as session: pairs = list((await session.execute(stmt)).all()) rows = [{**_dict(vector), "page": _dict(page)} for vector, page in pairs] if not rows: break for row in rows: yield row if len(rows) < batch: break offset += len(rows) async def get_index_revision(self, mapping_id: str) -> tuple[int, str | None, str]: async with self._sf() as session: state = await session.get(LlmWikiWikiSyncStateRow, mapping_id) if state is None: return 0, None, "not_synced" return int(state.index_revision or 0), state.embedding_fingerprint, state.state async def try_acquire_sync_lease(self, mapping_id: str, *, owner: str, sync_id: str, lease_seconds: int, fingerprint: str) -> bool: now = _now() async with self._sf() as session: try: await self._ensure_state(session, mapping_id) await session.commit() except IntegrityError: # Another worker created the one-per-KB state row between our # read and insert. Roll back this isolated bootstrap session; # the conditional UPDATE below is the actual lease election. await session.rollback() async with self._sf() as session: stmt = ( update(LlmWikiWikiSyncStateRow) .where( LlmWikiWikiSyncStateRow.knowledge_base_mapping_id == mapping_id, or_( LlmWikiWikiSyncStateRow.lease_owner.is_(None), LlmWikiWikiSyncStateRow.lease_expires_at.is_(None), LlmWikiWikiSyncStateRow.lease_expires_at < now, ), ) .values( state=case((LlmWikiWikiSyncStateRow.state == "rebuild_required", "rebuilding"), else_="running"), active_sync_id=sync_id, lease_owner=owner, lease_expires_at=now + timedelta(seconds=lease_seconds), last_started_at=now, last_error=None, updated_at=now, ) ) result = await session.execute(stmt) await session.commit() return bool(result.rowcount) async def renew_sync_lease(self, mapping_id: str, *, owner: str, lease_seconds: int) -> bool: now = _now() async with self._sf() as session: result = await session.execute( update(LlmWikiWikiSyncStateRow) .where(LlmWikiWikiSyncStateRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiSyncStateRow.lease_owner == owner) .values(lease_expires_at=now + timedelta(seconds=lease_seconds), updated_at=now) ) await session.commit() return bool(result.rowcount) async def update_sync_progress( self, mapping_id: str, *, owner: str, sync_id: str, remote_page_count: int, ) -> bool: async with self._sf() as session: result = await session.execute( update(LlmWikiWikiSyncStateRow) .where( LlmWikiWikiSyncStateRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiSyncStateRow.lease_owner == owner, LlmWikiWikiSyncStateRow.active_sync_id == sync_id, LlmWikiWikiSyncStateRow.lease_expires_at.is_not(None), LlmWikiWikiSyncStateRow.lease_expires_at >= _now(), ) .values(remote_page_count=max(0, remote_page_count), updated_at=_now()) ) await session.commit() return bool(result.rowcount) async def finish_sync( self, mapping_id: str, *, owner: str, sync_id: str, success: bool, remote_page_count: int, error: str | None = None, result: dict[str, Any] | None = None, ) -> bool: async with self._sf() as session: if not await self._fence_current_sync_lease( session, mapping_id, owner=owner, sync_id=sync_id, ): return False state = await session.get(LlmWikiWikiSyncStateRow, mapping_id) if state is None or state.lease_owner != owner or state.active_sync_id != sync_id: return False visible = [LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.is_remote_deleted.is_(False)] local_count = int((await session.execute(select(func.count()).select_from(LlmWikiWikiPageRow).where(*visible))).scalar_one()) ready_count = int((await session.execute(select(func.count()).select_from(LlmWikiWikiPageRow).where(*visible, LlmWikiWikiPageRow.index_status == "ready"))).scalar_one()) failed_count = int( ( await session.execute( select(func.count()) .select_from(LlmWikiWikiPageRow) .where( *visible, or_( LlmWikiWikiPageRow.index_status.in_(["failed", "stale"]), LlmWikiWikiPageRow.last_error.is_not(None), ), ) ) ).scalar_one() ) vector_count = int((await session.execute(select(func.count()).select_from(LlmWikiWikiVectorRow).where(LlmWikiWikiVectorRow.knowledge_base_mapping_id == mapping_id))).scalar_one()) now = _now() was_rebuilding = state.state == "rebuilding" state.state = "partial" if success and error else "idle" if success else "rebuild_required" if was_rebuilding else "failed" state.active_sync_id = None state.lease_owner = None state.lease_expires_at = None state.last_error = (error or "")[:4000] or None state.remote_page_count = remote_page_count state.local_page_count = local_count state.ready_page_count = ready_count state.failed_page_count = failed_count state.vector_count = vector_count if result is not None: state.last_sync_result = dict(result) if success and not error: state.last_completed_at = now if success: state.last_success_at = now state.updated_at = now await session.commit() return True async def list_index_status(self, mapping_ids: list[str] | None = None) -> list[dict[str, Any]]: stmt = select(LlmWikiWikiSyncStateRow) if mapping_ids: stmt = stmt.where(LlmWikiWikiSyncStateRow.knowledge_base_mapping_id.in_(mapping_ids)) stmt = stmt.order_by(LlmWikiWikiSyncStateRow.updated_at.desc()) async with self._sf() as session: return [_dict(row) for row in (await session.execute(stmt)).scalars()] async def list_page_index_status(self, mapping_id: str) -> list[dict[str, Any]]: vector_counts = ( select( LlmWikiWikiVectorRow.page_id.label("page_id"), func.count(LlmWikiWikiVectorRow.id).label("vector_count"), ) .where(LlmWikiWikiVectorRow.knowledge_base_mapping_id == mapping_id) .group_by(LlmWikiWikiVectorRow.page_id) .subquery() ) stmt = ( select(LlmWikiWikiPageRow, func.coalesce(vector_counts.c.vector_count, 0)) .outerjoin(vector_counts, vector_counts.c.page_id == LlmWikiWikiPageRow.id) .where( LlmWikiWikiPageRow.knowledge_base_mapping_id == mapping_id, LlmWikiWikiPageRow.is_remote_deleted.is_(False), ) .order_by(LlmWikiWikiPageRow.wiki_path, LlmWikiWikiPageRow.slug) ) async with self._sf() as session: rows = (await session.execute(stmt)).all() result: list[dict[str, Any]] = [] for row, vector_count in rows: page = _dict(row) result.append( { "id": page["id"], "slug": page["slug"], "title": page["title"], "page_type": page["page_type"], "status": page["status"], "index_status": page["index_status"], "vector_count": int(vector_count or 0), "indexed_at": page.get("indexed_at"), "last_error": page.get("last_error"), "last_seen_sync_id": page.get("last_seen_sync_id"), "updated_at": page.get("updated_at"), } ) return result async def mark_rebuild_required(self, mapping_ids: list[str], fingerprint: str) -> int: if not mapping_ids: return 0 now = _now() async with self._sf() as session: for mapping_id in mapping_ids: state = await self._ensure_state(session, mapping_id) state.state = "rebuild_required" state.embedding_fingerprint = fingerprint state.updated_at = now await session.execute(update(LlmWikiWikiPageRow).where(LlmWikiWikiPageRow.knowledge_base_mapping_id.in_(mapping_ids), LlmWikiWikiPageRow.is_remote_deleted.is_(False)).values(index_status="stale", updated_at=now)) await session.commit() return len(mapping_ids)