347 lines
14 KiB
Python
347 lines
14 KiB
Python
"""SQLAlchemy-backed generated-file library index."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import and_, delete, func, or_, select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.artifact_library.base import ArtifactLibraryStore
|
|
from deerflow.persistence.artifact_library.model import (
|
|
ArtifactLibraryItemRow,
|
|
ArtifactLibrarySyncStateRow,
|
|
)
|
|
|
|
_EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
|
|
|
|
|
|
def _artifact_id(user_id: str, item: dict[str, Any]) -> str:
|
|
identity = "\0".join((user_id, str(item.get("thread_id") or ""), str(item.get("path") or "")))
|
|
return hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _parse_datetime(value: object) -> datetime:
|
|
if isinstance(value, datetime):
|
|
parsed = value
|
|
else:
|
|
text = str(value or "").strip()
|
|
if not text:
|
|
return _EPOCH
|
|
try:
|
|
parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))
|
|
except ValueError:
|
|
return _EPOCH
|
|
if parsed.tzinfo is None:
|
|
return parsed.replace(tzinfo=UTC)
|
|
return parsed.astimezone(UTC)
|
|
|
|
|
|
def _iso(value: datetime | None) -> str | None:
|
|
return value.isoformat() if isinstance(value, datetime) else None
|
|
|
|
|
|
def _state_dict(row: ArtifactLibrarySyncStateRow | None) -> dict[str, Any]:
|
|
if row is None:
|
|
return {
|
|
"status": "idle",
|
|
"generation": None,
|
|
"started_at": None,
|
|
"completed_at": None,
|
|
"last_error": None,
|
|
}
|
|
return {
|
|
"status": row.status,
|
|
"generation": row.generation,
|
|
"started_at": _iso(row.started_at),
|
|
"completed_at": _iso(row.completed_at),
|
|
"last_error": row.last_error,
|
|
}
|
|
|
|
|
|
def _item_dict(row: ArtifactLibraryItemRow) -> dict[str, Any]:
|
|
return {
|
|
"thread_id": row.thread_id,
|
|
"thread_title": row.thread_title,
|
|
"agent_id": row.agent_id,
|
|
"agent_name": row.agent_name,
|
|
"thread_type": row.thread_type,
|
|
"source_type": row.source_type,
|
|
"strategy_project_id": row.strategy_project_id,
|
|
"strategy_name": row.strategy_name,
|
|
"strategy_stale": row.strategy_stale,
|
|
"strategy_source_kind": row.strategy_source_kind,
|
|
"download_url": row.download_url,
|
|
"path": row.path,
|
|
"name": row.name,
|
|
"size_bytes": row.size_bytes,
|
|
"mime_type": row.mime_type,
|
|
"modified_at": row.modified_at.isoformat(),
|
|
"file_kind": row.file_kind,
|
|
}
|
|
|
|
|
|
def _apply_item(
|
|
row: ArtifactLibraryItemRow,
|
|
*,
|
|
user_id: str,
|
|
generation: str,
|
|
item: dict[str, Any],
|
|
indexed_at: datetime,
|
|
) -> None:
|
|
row.user_id = user_id
|
|
row.thread_id = str(item.get("thread_id") or "")[:128]
|
|
row.thread_title = str(item.get("thread_title") or "")[:512]
|
|
row.agent_id = str(item.get("agent_id"))[:128] if item.get("agent_id") else None
|
|
row.agent_name = str(item.get("agent_name"))[:512] if item.get("agent_name") else None
|
|
row.thread_type = str(item.get("thread_type"))[:64] if item.get("thread_type") else None
|
|
row.source_type = str(item.get("source_type") or "general_qa")[:32]
|
|
row.strategy_project_id = str(item.get("strategy_project_id"))[:128] if item.get("strategy_project_id") else None
|
|
row.strategy_name = str(item.get("strategy_name"))[:512] if item.get("strategy_name") else None
|
|
row.strategy_stale = item.get("strategy_stale") is True
|
|
row.strategy_source_kind = str(item.get("strategy_source_kind"))[:32] if item.get("strategy_source_kind") else None
|
|
row.download_url = str(item.get("download_url")) if item.get("download_url") else None
|
|
row.path = str(item.get("path") or "")
|
|
row.name = str(item.get("name") or "")[:512]
|
|
row.size_bytes = max(0, int(item.get("size_bytes") or 0))
|
|
row.mime_type = str(item.get("mime_type"))[:255] if item.get("mime_type") else None
|
|
row.modified_at = _parse_datetime(item.get("modified_at"))
|
|
row.file_kind = str(item.get("file_kind") or "other")[:32]
|
|
row.sync_generation = generation
|
|
row.indexed_at = indexed_at
|
|
|
|
|
|
class SqlArtifactLibraryStore(ArtifactLibraryStore):
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
async def claim_sync(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
force: bool = False,
|
|
stale_after_seconds: float = 15.0,
|
|
lease_seconds: float = 300.0,
|
|
) -> str | None:
|
|
now = datetime.now(UTC)
|
|
generation = str(uuid4())
|
|
try:
|
|
async with self._sf() as session:
|
|
async with session.begin():
|
|
claim_filters = [
|
|
ArtifactLibrarySyncStateRow.user_id == user_id,
|
|
# An expired (or legacy missing) lease may be reclaimed.
|
|
# A live lease remains exclusive even for force=True.
|
|
or_(
|
|
ArtifactLibrarySyncStateRow.status != "syncing",
|
|
ArtifactLibrarySyncStateRow.lease_expires_at.is_(None),
|
|
ArtifactLibrarySyncStateRow.lease_expires_at <= now,
|
|
),
|
|
]
|
|
if not force:
|
|
stale_before = now - timedelta(seconds=stale_after_seconds)
|
|
claim_filters.append(
|
|
or_(
|
|
and_(
|
|
ArtifactLibrarySyncStateRow.status == "error",
|
|
or_(
|
|
ArtifactLibrarySyncStateRow.started_at.is_(None),
|
|
ArtifactLibrarySyncStateRow.started_at <= stale_before,
|
|
),
|
|
),
|
|
and_(
|
|
ArtifactLibrarySyncStateRow.status != "error",
|
|
or_(
|
|
ArtifactLibrarySyncStateRow.completed_at.is_(None),
|
|
ArtifactLibrarySyncStateRow.completed_at <= stale_before,
|
|
),
|
|
),
|
|
)
|
|
)
|
|
|
|
# A conditional UPDATE is the portable compare-and-swap.
|
|
# SQLite ignores SELECT ... FOR UPDATE, whereas this write
|
|
# serializes claimers there and also works on MySQL,
|
|
# PostgreSQL, and openGauss.
|
|
result = await session.execute(
|
|
update(ArtifactLibrarySyncStateRow)
|
|
.where(*claim_filters)
|
|
.values(
|
|
status="syncing",
|
|
generation=generation,
|
|
started_at=now,
|
|
lease_expires_at=now + timedelta(seconds=lease_seconds),
|
|
last_error=None,
|
|
)
|
|
)
|
|
if result.rowcount == 1:
|
|
return generation
|
|
|
|
row = await session.get(ArtifactLibrarySyncStateRow, user_id)
|
|
if row is not None:
|
|
return None
|
|
session.add(
|
|
ArtifactLibrarySyncStateRow(
|
|
user_id=user_id,
|
|
status="syncing",
|
|
generation=generation,
|
|
started_at=now,
|
|
lease_expires_at=now + timedelta(seconds=lease_seconds),
|
|
last_error=None,
|
|
)
|
|
)
|
|
return generation
|
|
except IntegrityError:
|
|
# Two workers can race while inserting a user's first state row.
|
|
# The winner owns the lease; the loser simply serves the same
|
|
# persisted index without starting a duplicate filesystem scan.
|
|
async with self._sf() as session:
|
|
if await session.get(ArtifactLibrarySyncStateRow, user_id) is not None:
|
|
return None
|
|
raise
|
|
|
|
async def finish_sync(
|
|
self,
|
|
user_id: str,
|
|
generation: str,
|
|
items: list[dict[str, Any]],
|
|
) -> bool:
|
|
now = datetime.now(UTC)
|
|
deduped = {_artifact_id(user_id, item): item for item in items}
|
|
async with self._sf() as session:
|
|
async with session.begin():
|
|
# Write first so SQLite also serializes this publisher against
|
|
# a lease takeover. Requiring status=syncing makes completion
|
|
# one-shot even if the same callback is accidentally invoked
|
|
# twice with an old generation token.
|
|
fence = await session.execute(
|
|
update(ArtifactLibrarySyncStateRow)
|
|
.where(
|
|
ArtifactLibrarySyncStateRow.user_id == user_id,
|
|
ArtifactLibrarySyncStateRow.generation == generation,
|
|
ArtifactLibrarySyncStateRow.status == "syncing",
|
|
)
|
|
.values(lease_expires_at=now + timedelta(seconds=300))
|
|
)
|
|
if fence.rowcount != 1:
|
|
return False
|
|
|
|
existing_result = await session.execute(select(ArtifactLibraryItemRow).where(ArtifactLibraryItemRow.user_id == user_id))
|
|
existing = {row.artifact_id: row for row in existing_result.scalars()}
|
|
for artifact_id, item in deduped.items():
|
|
row = existing.get(artifact_id)
|
|
if row is None:
|
|
row = ArtifactLibraryItemRow(
|
|
artifact_id=artifact_id,
|
|
user_id=user_id,
|
|
thread_id="",
|
|
thread_title="",
|
|
source_type="general_qa",
|
|
path="",
|
|
name="",
|
|
size_bytes=0,
|
|
modified_at=_EPOCH,
|
|
file_kind="other",
|
|
sync_generation=generation,
|
|
indexed_at=now,
|
|
)
|
|
session.add(row)
|
|
_apply_item(
|
|
row,
|
|
user_id=user_id,
|
|
generation=generation,
|
|
item=item,
|
|
indexed_at=now,
|
|
)
|
|
await session.flush()
|
|
await session.execute(
|
|
delete(ArtifactLibraryItemRow).where(
|
|
ArtifactLibraryItemRow.user_id == user_id,
|
|
ArtifactLibraryItemRow.sync_generation != generation,
|
|
)
|
|
)
|
|
await session.execute(
|
|
update(ArtifactLibrarySyncStateRow)
|
|
.where(
|
|
ArtifactLibrarySyncStateRow.user_id == user_id,
|
|
ArtifactLibrarySyncStateRow.generation == generation,
|
|
ArtifactLibrarySyncStateRow.status == "syncing",
|
|
)
|
|
.values(
|
|
status="ready",
|
|
completed_at=now,
|
|
lease_expires_at=None,
|
|
last_error=None,
|
|
)
|
|
)
|
|
return True
|
|
|
|
async def fail_sync(self, user_id: str, generation: str, error: str) -> bool:
|
|
async with self._sf() as session:
|
|
async with session.begin():
|
|
result = await session.execute(
|
|
update(ArtifactLibrarySyncStateRow)
|
|
.where(
|
|
ArtifactLibrarySyncStateRow.user_id == user_id,
|
|
ArtifactLibrarySyncStateRow.generation == generation,
|
|
ArtifactLibrarySyncStateRow.status == "syncing",
|
|
)
|
|
.values(
|
|
status="error",
|
|
lease_expires_at=None,
|
|
last_error=(error or "Artifact library refresh failed")[:1024],
|
|
)
|
|
)
|
|
return result.rowcount == 1
|
|
|
|
async def get_sync_state(self, user_id: str) -> dict[str, Any]:
|
|
async with self._sf() as session:
|
|
row = await session.get(ArtifactLibrarySyncStateRow, user_id)
|
|
return _state_dict(row)
|
|
|
|
async def list_items(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
query: str | None = None,
|
|
file_kind: str = "all",
|
|
source_type: str = "all",
|
|
offset: int = 0,
|
|
limit: int = 50,
|
|
) -> tuple[int, list[dict[str, Any]]]:
|
|
filters = [ArtifactLibraryItemRow.user_id == user_id]
|
|
if file_kind != "all":
|
|
filters.append(ArtifactLibraryItemRow.file_kind == file_kind)
|
|
if source_type != "all":
|
|
filters.append(ArtifactLibraryItemRow.source_type == source_type)
|
|
normalized = (query or "").strip().casefold()
|
|
if normalized:
|
|
searchable = (
|
|
ArtifactLibraryItemRow.name,
|
|
ArtifactLibraryItemRow.thread_title,
|
|
ArtifactLibraryItemRow.agent_name,
|
|
ArtifactLibraryItemRow.strategy_name,
|
|
ArtifactLibraryItemRow.strategy_project_id,
|
|
)
|
|
filters.append(or_(*(func.lower(func.coalesce(column, "")).contains(normalized, autoescape=True) for column in searchable)))
|
|
|
|
count_stmt = select(func.count()).select_from(ArtifactLibraryItemRow).where(*filters)
|
|
page_stmt = (
|
|
select(ArtifactLibraryItemRow)
|
|
.where(*filters)
|
|
.order_by(
|
|
ArtifactLibraryItemRow.modified_at.desc(),
|
|
ArtifactLibraryItemRow.artifact_id.desc(),
|
|
)
|
|
.offset(offset)
|
|
.limit(limit)
|
|
)
|
|
async with self._sf() as session:
|
|
total = int((await session.execute(count_stmt)).scalar_one())
|
|
rows = (await session.execute(page_stmt)).scalars().all()
|
|
return total, [_item_dict(row) for row in rows]
|