feat: add checkpoint postgres migration admin
This commit is contained in:
1 parent
eaa061dd06
commit
9471e72171
10 files changed
+962
-33
No files matched your search
@@ -4,9 +4,10 @@ import {
|
||||
BellIcon,
|
||||
BotIcon,
|
||||
BrainIcon,
|
||||
CalendarClockIcon,
|
||||
ChevronsUpDown,
|
||||
GaugeIcon,
|
||||
CalendarClockIcon,
|
||||
ChevronsUpDown,
|
||||
DatabaseZapIcon,
|
||||
GaugeIcon,
|
||||
HistoryIcon,
|
||||
MessageSquareCodeIcon,
|
||||
MessageSquarePlusIcon,
|
||||
@@ -198,6 +199,16 @@ export function WorkspaceNavMenu() {
|
||||
用户排行榜
|
||||
</DropdownMenuItem>
|
||||
)}
|
||||
{isAdmin && (
|
||||
<DropdownMenuItem
|
||||
onClick={() =>
|
||||
navigate("/page/workspace/settings/checkpoint-migration")
|
||||
}
|
||||
>
|
||||
<DatabaseZapIcon />
|
||||
Checkpoint 迁移
|
||||
</DropdownMenuItem>
|
||||
)}
|
||||
{isAdmin && (
|
||||
<DropdownMenuItem
|
||||
onClick={() => navigate("/page/workspace/admin")}
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
"use client";
|
||||
|
||||
import { DatabaseZap, Loader2Icon, RefreshCwIcon } from "lucide-react";
|
||||
import { useMutation, useQuery } from "@tanstack/react-query";
|
||||
import { useMemo, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { apiFetch } from "@/core/api";
|
||||
import { getBackendBaseURL } from "@/core/config";
|
||||
|
||||
interface CheckpointMigrationSources {
|
||||
source_paths: string[];
|
||||
}
|
||||
|
||||
interface CheckpointMigrationTask {
|
||||
task_id: string;
|
||||
status: "queued" | "running" | "completed" | "failed";
|
||||
start_at: string;
|
||||
end_at: string;
|
||||
source_paths: string[];
|
||||
total_threads: number;
|
||||
copied_threads: number;
|
||||
skipped_threads: number;
|
||||
failed_threads: number;
|
||||
skipped_ranges: number;
|
||||
message: string;
|
||||
error?: string | null;
|
||||
created_at: string;
|
||||
updated_at: string;
|
||||
}
|
||||
|
||||
async function loadSources(): Promise<CheckpointMigrationSources> {
|
||||
const response = await apiFetch(
|
||||
`${getBackendBaseURL()}/api/admin/checkpoint-migration/sources`,
|
||||
);
|
||||
if (!response.ok) throw new Error("加载 SQLite 来源失败");
|
||||
return (await response.json()) as CheckpointMigrationSources;
|
||||
}
|
||||
|
||||
async function startMigration(body: {
|
||||
start_at: string;
|
||||
end_at: string;
|
||||
}): Promise<CheckpointMigrationTask> {
|
||||
const response = await apiFetch(
|
||||
`${getBackendBaseURL()}/api/admin/checkpoint-migration/tasks`,
|
||||
{
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(body),
|
||||
},
|
||||
);
|
||||
if (!response.ok) {
|
||||
const detail = await response.json().catch(() => ({ detail: "启动迁移失败" }));
|
||||
throw new Error(detail.detail ?? "启动迁移失败");
|
||||
}
|
||||
return (await response.json()) as CheckpointMigrationTask;
|
||||
}
|
||||
|
||||
async function loadTask(taskId: string): Promise<CheckpointMigrationTask> {
|
||||
const response = await apiFetch(
|
||||
`${getBackendBaseURL()}/api/admin/checkpoint-migration/tasks/${taskId}`,
|
||||
);
|
||||
if (!response.ok) throw new Error("加载迁移任务失败");
|
||||
return (await response.json()) as CheckpointMigrationTask;
|
||||
}
|
||||
|
||||
function toLocalInputValue(date: Date): string {
|
||||
const pad = (n: number) => String(n).padStart(2, "0");
|
||||
return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}T${pad(date.getHours())}:${pad(date.getMinutes())}`;
|
||||
}
|
||||
|
||||
function statusVariant(status?: string) {
|
||||
if (status === "failed") return "destructive" as const;
|
||||
if (status === "completed") return "default" as const;
|
||||
if (status === "queued" || status === "running") return "secondary" as const;
|
||||
return "outline" as const;
|
||||
}
|
||||
|
||||
export default function CheckpointMigrationPage() {
|
||||
const now = useMemo(() => new Date(), []);
|
||||
const weekAgo = useMemo(
|
||||
() => new Date(now.getTime() - 7 * 24 * 60 * 60 * 1000),
|
||||
[now],
|
||||
);
|
||||
const [startAt, setStartAt] = useState(toLocalInputValue(weekAgo));
|
||||
const [endAt, setEndAt] = useState(toLocalInputValue(now));
|
||||
const [taskId, setTaskId] = useState<string | null>(null);
|
||||
|
||||
const { data: sources, isLoading: sourcesLoading, refetch } = useQuery({
|
||||
queryKey: ["checkpointMigration", "sources"],
|
||||
queryFn: loadSources,
|
||||
staleTime: 60_000,
|
||||
refetchOnWindowFocus: false,
|
||||
});
|
||||
const startMutation = useMutation({ mutationFn: startMigration });
|
||||
const { data: task } = useQuery({
|
||||
queryKey: ["checkpointMigration", "task", taskId],
|
||||
queryFn: () => loadTask(taskId as string),
|
||||
enabled: !!taskId,
|
||||
refetchInterval: (query) => {
|
||||
const status = query.state.data?.status;
|
||||
return status === "queued" || status === "running" ? 2000 : false;
|
||||
},
|
||||
refetchOnWindowFocus: false,
|
||||
});
|
||||
|
||||
const running =
|
||||
startMutation.isPending ||
|
||||
task?.status === "queued" ||
|
||||
task?.status === "running";
|
||||
|
||||
async function handleStart() {
|
||||
const start = new Date(startAt);
|
||||
const end = new Date(endAt);
|
||||
if (Number.isNaN(start.getTime()) || Number.isNaN(end.getTime())) {
|
||||
toast.error("请选择有效的开始和结束时间");
|
||||
return;
|
||||
}
|
||||
if (end <= start) {
|
||||
toast.error("结束时间必须晚于开始时间");
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const created = await startMutation.mutateAsync({
|
||||
start_at: start.toISOString(),
|
||||
end_at: end.toISOString(),
|
||||
});
|
||||
setTaskId(created.task_id);
|
||||
toast.success("迁移任务已启动");
|
||||
} catch (err) {
|
||||
toast.error(err instanceof Error ? err.message : "启动迁移失败");
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mx-auto flex w-full max-w-3xl flex-col gap-6 px-6 py-10">
|
||||
<div>
|
||||
<h1 className="text-xl font-semibold">Checkpoint 迁移</h1>
|
||||
<p className="text-muted-foreground mt-1 text-sm">
|
||||
按时间段把 SQLite 旧聊天 checkpoint 同步到 Postgres。同一来源、同一时间段完成后再次执行会跳过。
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<section className="rounded-lg border bg-white p-5 dark:border-[#213D5E] dark:bg-[#02244C]">
|
||||
<div className="mb-4 flex items-center justify-between">
|
||||
<div className="flex items-center gap-2 text-sm font-medium">
|
||||
<DatabaseZap className="size-4" />
|
||||
SQLite 来源
|
||||
</div>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
onClick={() => refetch()}
|
||||
disabled={sourcesLoading}
|
||||
title="刷新"
|
||||
>
|
||||
<RefreshCwIcon className="size-4" />
|
||||
</Button>
|
||||
</div>
|
||||
{sourcesLoading ? (
|
||||
<div className="text-muted-foreground flex items-center gap-2 text-sm">
|
||||
<Loader2Icon className="size-4 animate-spin" />
|
||||
加载中
|
||||
</div>
|
||||
) : (sources?.source_paths.length ?? 0) > 0 ? (
|
||||
<div className="space-y-2">
|
||||
{sources?.source_paths.map((path) => (
|
||||
<div
|
||||
key={path}
|
||||
className="rounded border bg-muted/40 px-3 py-2 font-mono text-xs break-all"
|
||||
>
|
||||
{path}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
) : (
|
||||
<div className="text-muted-foreground rounded border border-dashed px-3 py-6 text-center text-sm">
|
||||
未配置 SQLite 来源路径
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
|
||||
<section className="rounded-lg border bg-white p-5 dark:border-[#213D5E] dark:bg-[#02244C]">
|
||||
<div className="grid gap-4 sm:grid-cols-2">
|
||||
<label className="flex flex-col gap-1.5 text-sm">
|
||||
<span className="font-medium">开始时间</span>
|
||||
<Input
|
||||
type="datetime-local"
|
||||
value={startAt}
|
||||
onChange={(e) => setStartAt(e.target.value)}
|
||||
/>
|
||||
</label>
|
||||
<label className="flex flex-col gap-1.5 text-sm">
|
||||
<span className="font-medium">结束时间</span>
|
||||
<Input
|
||||
type="datetime-local"
|
||||
value={endAt}
|
||||
onChange={(e) => setEndAt(e.target.value)}
|
||||
/>
|
||||
</label>
|
||||
</div>
|
||||
<div className="mt-5 flex justify-end">
|
||||
<Button
|
||||
onClick={handleStart}
|
||||
disabled={running || !(sources?.source_paths.length ?? 0)}
|
||||
>
|
||||
{running && <Loader2Icon className="size-4 animate-spin" />}
|
||||
开始同步
|
||||
</Button>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
{task && (
|
||||
<section className="rounded-lg border bg-white p-5 dark:border-[#213D5E] dark:bg-[#02244C]">
|
||||
<div className="mb-4 flex items-center justify-between">
|
||||
<div className="text-sm font-medium">任务状态</div>
|
||||
<Badge variant={statusVariant(task.status)}>{task.status}</Badge>
|
||||
</div>
|
||||
<div className="grid gap-3 text-sm sm:grid-cols-4">
|
||||
<div className="rounded border p-3">
|
||||
<div className="text-muted-foreground text-xs">总线程</div>
|
||||
<div className="mt-1 text-lg font-semibold tabular-nums">
|
||||
{task.total_threads}
|
||||
</div>
|
||||
</div>
|
||||
<div className="rounded border p-3">
|
||||
<div className="text-muted-foreground text-xs">已复制</div>
|
||||
<div className="mt-1 text-lg font-semibold tabular-nums">
|
||||
{task.copied_threads}
|
||||
</div>
|
||||
</div>
|
||||
<div className="rounded border p-3">
|
||||
<div className="text-muted-foreground text-xs">已跳过</div>
|
||||
<div className="mt-1 text-lg font-semibold tabular-nums">
|
||||
{task.skipped_threads}
|
||||
</div>
|
||||
</div>
|
||||
<div className="rounded border p-3">
|
||||
<div className="text-muted-foreground text-xs">跳过时间段</div>
|
||||
<div className="mt-1 text-lg font-semibold tabular-nums">
|
||||
{task.skipped_ranges}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{(task.message || task.error) && (
|
||||
<div className="mt-4 rounded border bg-muted/40 px-3 py-2 text-sm">
|
||||
{task.error || task.message}
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -25,6 +25,7 @@ import AdminSkillsPage from './AdminSkillsPage';
|
||||
import ScheduledTaskRunDetailPage from './ScheduledTaskRunDetailPage';
|
||||
import ScheduledTasksPage from './ScheduledTasksPage';
|
||||
import LlmMetricsPage from './LlmMetricsPage';
|
||||
import CheckpointMigrationPage from './CheckpointMigrationPage';
|
||||
import MemoryPage from './MemoryPage';
|
||||
import PromptPrefixSettingsPage from './PromptPrefixSettingsPage';
|
||||
import PrivateSquareSettingsPage from './PrivateSquareSettingsPage';
|
||||
@@ -476,6 +477,18 @@ export default function WorkspaceRoutes({
|
||||
)
|
||||
}
|
||||
/>
|
||||
<Route
|
||||
path="settings/checkpoint-migration"
|
||||
element={
|
||||
getStoredAuth()?.system_role === 'admin' ? (
|
||||
<WorkspaceLayout {...layoutProps}>
|
||||
<CheckpointMigrationPage />
|
||||
</WorkspaceLayout>
|
||||
) : (
|
||||
<Navigate to="/page/workspace/chats/new" replace />
|
||||
)
|
||||
}
|
||||
/>
|
||||
</Routes>
|
||||
</RouteBaseProvider>
|
||||
);
|
||||
|
||||
@@ -30,8 +30,9 @@ from app.gateway.routers import (
|
||||
artifacts,
|
||||
assistants_compat,
|
||||
auth,
|
||||
channels,
|
||||
feedback,
|
||||
channels,
|
||||
checkpoint_migration,
|
||||
feedback,
|
||||
html_page_favorites,
|
||||
intent,
|
||||
llm_metrics,
|
||||
@@ -905,7 +906,10 @@ This gateway provides custom endpoints for models, MCP configuration, skills, an
|
||||
app.include_router(tool_metrics.router)
|
||||
|
||||
# Channels API is mounted at /api/channels
|
||||
app.include_router(channels.router)
|
||||
app.include_router(channels.router)
|
||||
|
||||
# Admin checkpoint migration API: SQLite legacy checkpoints -> Postgres.
|
||||
app.include_router(checkpoint_migration.router)
|
||||
|
||||
# Assistants compatibility API (LangGraph Platform stub)
|
||||
app.include_router(assistants_compat.router)
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Admin checkpoint migration endpoints.
|
||||
|
||||
These endpoints support the SQLite -> Postgres transition for LangGraph
|
||||
checkpoints. The live checkpointer should be Postgres with SQLite dual-read
|
||||
enabled; this router only starts manual, bounded migrations for selected time
|
||||
ranges.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.auth.disabled_mode import is_auth_disabled
|
||||
from app.gateway.deps import get_current_user_from_request
|
||||
from deerflow.config.runtime_paths import runtime_home
|
||||
from deerflow.runtime.checkpointer.checkpoint_migration import (
|
||||
SQLiteCheckpointSource,
|
||||
copy_thread_state,
|
||||
sqlite_recent_threads,
|
||||
target_has_thread,
|
||||
)
|
||||
from deerflow.runtime.store._sqlite_utils import resolve_sqlite_conn_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/admin/checkpoint-migration", tags=["checkpoint-migration"])
|
||||
|
||||
_TASKS: dict[str, dict[str, Any]] = {}
|
||||
|
||||
|
||||
class MigrationStartRequest(BaseModel):
|
||||
start_at: str = Field(description="Inclusive start datetime, ISO string.")
|
||||
end_at: str = Field(description="Inclusive end datetime, ISO string.")
|
||||
source_paths: list[str] | None = Field(default=None, description="Optional SQLite sources. Defaults to configured dual-read paths.")
|
||||
|
||||
|
||||
class MigrationTaskResponse(BaseModel):
|
||||
task_id: str
|
||||
status: str
|
||||
start_at: str
|
||||
end_at: str
|
||||
source_paths: list[str]
|
||||
total_threads: int = 0
|
||||
copied_threads: int = 0
|
||||
skipped_threads: int = 0
|
||||
failed_threads: int = 0
|
||||
skipped_ranges: int = 0
|
||||
message: str = ""
|
||||
error: str | None = None
|
||||
created_at: str
|
||||
updated_at: str
|
||||
|
||||
|
||||
async def _require_admin(request: Request) -> None:
|
||||
if is_auth_disabled():
|
||||
return
|
||||
user = await get_current_user_from_request(request)
|
||||
if getattr(user, "system_role", None) != "admin":
|
||||
raise HTTPException(status_code=403, detail="Checkpoint migration requires admin role")
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _parse_dt(value: str) -> datetime:
|
||||
try:
|
||||
normalized = value.replace("Z", "+00:00")
|
||||
dt = datetime.fromisoformat(normalized)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid datetime: {value}") from exc
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt
|
||||
|
||||
|
||||
def _sqlite_datetime(dt: datetime) -> str:
|
||||
# Existing SQLite rows store naive wall-clock strings like
|
||||
# "2026-05-18 06:41:10.989427"; match that format for comparisons.
|
||||
return dt.astimezone(timezone.utc).replace(tzinfo=None).isoformat(sep=" ")
|
||||
|
||||
|
||||
def _configured_sources(request: Request) -> list[str]:
|
||||
config = getattr(request.app.state, "config", None)
|
||||
cp = getattr(config, "checkpointer", None)
|
||||
paths = list(getattr(cp, "dual_read_sqlite_paths", []) or [])
|
||||
raw_env = os.getenv("DEER_FLOW_CHECKPOINT_SQLITE_PATHS", "")
|
||||
paths.extend(part.strip() for part in raw_env.split(",") if part.strip())
|
||||
resolved: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for path in paths:
|
||||
conn_str = resolve_sqlite_conn_str(path)
|
||||
if conn_str == ":memory:" or conn_str.startswith("file:"):
|
||||
continue
|
||||
if conn_str not in seen:
|
||||
resolved.append(conn_str)
|
||||
seen.add(conn_str)
|
||||
return resolved
|
||||
|
||||
|
||||
def _target_checkpointer(request: Request):
|
||||
checkpointer = getattr(request.app.state, "checkpointer", None)
|
||||
if checkpointer is None:
|
||||
raise HTTPException(status_code=503, detail="Checkpointer is not available")
|
||||
return getattr(checkpointer, "primary", checkpointer)
|
||||
|
||||
|
||||
def _state_path() -> Path:
|
||||
return runtime_home() / "checkpoint_migration_state.json"
|
||||
|
||||
|
||||
def _load_completed_ranges() -> set[str]:
|
||||
try:
|
||||
payload = json.loads(_state_path().read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return set()
|
||||
except Exception:
|
||||
logger.warning("Ignoring unreadable checkpoint migration state file", exc_info=True)
|
||||
return set()
|
||||
return {str(item) for item in payload.get("completed_ranges", [])}
|
||||
|
||||
|
||||
def _save_completed_ranges(items: set[str]) -> None:
|
||||
path = _state_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(".tmp")
|
||||
tmp.write_text(json.dumps({"completed_ranges": sorted(items)}, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
tmp.replace(path)
|
||||
|
||||
|
||||
def _range_key(source: str, start_at: str, end_at: str) -> str:
|
||||
raw = json.dumps({"source": str(Path(source).resolve()), "start_at": start_at, "end_at": end_at}, sort_keys=True)
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
async def _open_sqlite_source(path: str) -> SQLiteCheckpointSource:
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
ctx = AsyncSqliteSaver.from_conn_string(path)
|
||||
saver = await ctx.__aenter__()
|
||||
await saver.setup()
|
||||
return SQLiteCheckpointSource(
|
||||
path=path,
|
||||
saver=saver,
|
||||
context=ctx,
|
||||
progress_path=Path(path).with_suffix(Path(path).suffix + ".pg-migration.json"),
|
||||
)
|
||||
|
||||
|
||||
def _task_response(task: dict[str, Any]) -> MigrationTaskResponse:
|
||||
return MigrationTaskResponse(**task)
|
||||
|
||||
|
||||
async def _run_task(task_id: str, target, start_at: datetime, end_at: datetime) -> None:
|
||||
task = _TASKS[task_id]
|
||||
completed = await asyncio.to_thread(_load_completed_ranges)
|
||||
start_sql = _sqlite_datetime(start_at)
|
||||
end_sql = _sqlite_datetime(end_at)
|
||||
range_start = start_at.isoformat()
|
||||
range_end = end_at.isoformat()
|
||||
|
||||
try:
|
||||
task["status"] = "running"
|
||||
task["updated_at"] = _now_iso()
|
||||
for source_path in task["source_paths"]:
|
||||
if not Path(source_path).exists():
|
||||
task["message"] = f"SQLite source not found: {source_path}"
|
||||
task["skipped_ranges"] += 1
|
||||
continue
|
||||
|
||||
key = _range_key(source_path, range_start, range_end)
|
||||
if key in completed:
|
||||
task["skipped_ranges"] += 1
|
||||
task["message"] = "Selected range was already migrated; skipped"
|
||||
task["updated_at"] = _now_iso()
|
||||
continue
|
||||
|
||||
thread_ids = await asyncio.to_thread(sqlite_recent_threads, source_path, since=start_sql, until=end_sql)
|
||||
if not thread_ids:
|
||||
task["message"] = f"No threads found in selected range for {Path(source_path).name}"
|
||||
completed.add(key)
|
||||
await asyncio.to_thread(_save_completed_ranges, completed)
|
||||
task["updated_at"] = _now_iso()
|
||||
continue
|
||||
|
||||
task["total_threads"] += len(thread_ids)
|
||||
source = await _open_sqlite_source(source_path)
|
||||
try:
|
||||
for thread_id in thread_ids:
|
||||
if await target_has_thread(target, thread_id):
|
||||
task["skipped_threads"] += 1
|
||||
else:
|
||||
copied = await copy_thread_state(source.saver, target, thread_id)
|
||||
if copied:
|
||||
task["copied_threads"] += 1
|
||||
else:
|
||||
task["skipped_threads"] += 1
|
||||
task["updated_at"] = _now_iso()
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
await source.close()
|
||||
|
||||
completed.add(key)
|
||||
await asyncio.to_thread(_save_completed_ranges, completed)
|
||||
|
||||
task["status"] = "completed"
|
||||
task["message"] = task["message"] or "Migration completed"
|
||||
except Exception as exc:
|
||||
logger.exception("Checkpoint migration task failed: %s", task_id)
|
||||
task["status"] = "failed"
|
||||
task["error"] = str(exc)
|
||||
finally:
|
||||
task["updated_at"] = _now_iso()
|
||||
|
||||
|
||||
@router.post("/tasks", response_model=MigrationTaskResponse)
|
||||
async def start_migration(body: MigrationStartRequest, request: Request):
|
||||
await _require_admin(request)
|
||||
start_at = _parse_dt(body.start_at)
|
||||
end_at = _parse_dt(body.end_at)
|
||||
if end_at <= start_at:
|
||||
raise HTTPException(status_code=400, detail="end_at must be later than start_at")
|
||||
|
||||
sources = [resolve_sqlite_conn_str(p) for p in (body.source_paths or _configured_sources(request))]
|
||||
sources = [p for p in dict.fromkeys(sources) if p != ":memory:" and not p.startswith("file:")]
|
||||
if not sources:
|
||||
raise HTTPException(status_code=400, detail="No SQLite source paths configured")
|
||||
|
||||
target = _target_checkpointer(request)
|
||||
task_id = hashlib.sha256(f"{_now_iso()}:{sources}:{body.start_at}:{body.end_at}".encode("utf-8")).hexdigest()[:16]
|
||||
now = _now_iso()
|
||||
_TASKS[task_id] = {
|
||||
"task_id": task_id,
|
||||
"status": "queued",
|
||||
"start_at": start_at.isoformat(),
|
||||
"end_at": end_at.isoformat(),
|
||||
"source_paths": sources,
|
||||
"total_threads": 0,
|
||||
"copied_threads": 0,
|
||||
"skipped_threads": 0,
|
||||
"failed_threads": 0,
|
||||
"skipped_ranges": 0,
|
||||
"message": "",
|
||||
"error": None,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
asyncio.create_task(_run_task(task_id, target, start_at, end_at))
|
||||
return _task_response(_TASKS[task_id])
|
||||
|
||||
|
||||
@router.get("/tasks/{task_id}", response_model=MigrationTaskResponse)
|
||||
async def get_migration_task(task_id: str, request: Request):
|
||||
await _require_admin(request)
|
||||
task = _TASKS.get(task_id)
|
||||
if task is None:
|
||||
raise HTTPException(status_code=404, detail="Migration task not found")
|
||||
return _task_response(task)
|
||||
|
||||
|
||||
@router.get("/sources")
|
||||
async def list_sources(request: Request):
|
||||
await _require_admin(request)
|
||||
return {"source_paths": _configured_sources(request)}
|
||||
@@ -219,8 +219,12 @@ database:
|
||||
echo_sql: false
|
||||
|
||||
checkpointer:
|
||||
type: sqlite
|
||||
connection_string: .deer-flow/data/deerflow.db
|
||||
type: postgres
|
||||
connection_string: postgresql://postgres:postgres123@47.88.25.99:8432/deerflow
|
||||
dual_read_enabled: true
|
||||
dual_read_sqlite_paths:
|
||||
- .deer-flow/data/deerflow.db
|
||||
- ../data/.deer-flow/data/checkpoints.db
|
||||
|
||||
# 记忆子系统:跨会话持久化用户画像与 agent 私有记忆。
|
||||
# 三层覆盖优先级(高->低):
|
||||
@@ -301,4 +305,4 @@ ai_writing:
|
||||
log_system:
|
||||
enabled: true
|
||||
url: "https://your-log-system.example.com"
|
||||
label: "日志系统"
|
||||
label: "日志系统"
|
||||
+20
-6
@@ -2,7 +2,7 @@
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
CheckpointerType = Literal["memory", "sqlite", "postgres"]
|
||||
|
||||
@@ -16,13 +16,27 @@ class CheckpointerConfig(BaseModel):
|
||||
"'sqlite' persists to a local file (requires langgraph-checkpoint-sqlite). "
|
||||
"'postgres' persists to PostgreSQL (requires langgraph-checkpoint-postgres)."
|
||||
)
|
||||
connection_string: str | None = Field(
|
||||
default=None,
|
||||
description="Connection string for sqlite (file path) or postgres (DSN). "
|
||||
connection_string: str | None = Field(
|
||||
default=None,
|
||||
description="Connection string for sqlite (file path) or postgres (DSN). "
|
||||
"Required for sqlite and postgres types. "
|
||||
"For sqlite, use a file path like '.deer-flow/checkpoints.db' or ':memory:' for in-memory. "
|
||||
"For postgres, use a DSN like 'postgresql://user:pass@localhost:5432/db'.",
|
||||
)
|
||||
"For postgres, use a DSN like 'postgresql://user:pass@localhost:5432/db'.",
|
||||
)
|
||||
dual_read_enabled: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"When type='postgres', read from Postgres first and fall back to legacy "
|
||||
"SQLite checkpoint files when Postgres does not have the requested data."
|
||||
),
|
||||
)
|
||||
dual_read_sqlite_paths: list[str] = Field(
|
||||
default_factory=list,
|
||||
description=(
|
||||
"Legacy SQLite checkpoint database paths used for read fallback. "
|
||||
"Plain paths are resolved the same way as SQLite connection_string."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# Global configuration instance — None means no checkpointer is configured.
|
||||
|
||||
+86
-16
@@ -17,10 +17,12 @@ For sync usage see :mod:`deerflow.runtime.checkpointer.provider`.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from langgraph.types import Checkpointer
|
||||
|
||||
@@ -30,9 +32,37 @@ from deerflow.runtime.checkpointer.provider import (
|
||||
POSTGRES_INSTALL,
|
||||
SQLITE_INSTALL,
|
||||
)
|
||||
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
||||
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _csv_env(name: str) -> list[str]:
|
||||
raw = os.getenv(name, "")
|
||||
return [part.strip() for part in raw.split(",") if part.strip()]
|
||||
|
||||
|
||||
def _dual_read_enabled(config) -> bool:
|
||||
if getattr(config, "dual_read_enabled", False):
|
||||
return True
|
||||
return os.getenv("DEER_FLOW_CHECKPOINT_DUAL_READ", "").strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _dual_read_paths(config) -> list[str]:
|
||||
paths = list(getattr(config, "dual_read_sqlite_paths", []) or [])
|
||||
paths.extend(_csv_env("DEER_FLOW_CHECKPOINT_SQLITE_PATHS"))
|
||||
resolved: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for path in paths:
|
||||
conn_str = resolve_sqlite_conn_str(path)
|
||||
if conn_str == ":memory:" or conn_str.startswith("file:"):
|
||||
logger.warning("Skipping unsupported checkpoint migration SQLite path: %s", path)
|
||||
continue
|
||||
if conn_str not in seen:
|
||||
resolved.append(conn_str)
|
||||
seen.add(conn_str)
|
||||
return resolved
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async factory
|
||||
@@ -61,19 +91,59 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
|
||||
yield saver
|
||||
return
|
||||
|
||||
if config.type == "postgres":
|
||||
try:
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
if config.type == "postgres":
|
||||
try:
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
except ImportError as exc:
|
||||
raise ImportError(POSTGRES_INSTALL) from exc
|
||||
|
||||
if not config.connection_string:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
async with AsyncPostgresSaver.from_conn_string(config.connection_string) as saver:
|
||||
await saver.setup()
|
||||
yield saver
|
||||
return
|
||||
if not config.connection_string:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
async with AsyncPostgresSaver.from_conn_string(config.connection_string) as saver:
|
||||
await saver.setup()
|
||||
if not _dual_read_enabled(config):
|
||||
yield saver
|
||||
return
|
||||
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
from deerflow.runtime.checkpointer.checkpoint_migration import (
|
||||
DualReadCheckpointer,
|
||||
SQLiteCheckpointSource,
|
||||
)
|
||||
|
||||
sources: list[SQLiteCheckpointSource] = []
|
||||
for sqlite_path in _dual_read_paths(config):
|
||||
if not Path(sqlite_path).exists():
|
||||
logger.info("Checkpoint dual-read source does not exist, skipping: %s", sqlite_path)
|
||||
continue
|
||||
ctx = AsyncSqliteSaver.from_conn_string(sqlite_path)
|
||||
source_saver = await ctx.__aenter__()
|
||||
await source_saver.setup()
|
||||
sources.append(
|
||||
SQLiteCheckpointSource(
|
||||
path=sqlite_path,
|
||||
saver=source_saver,
|
||||
context=ctx,
|
||||
progress_path=Path(sqlite_path).with_suffix(Path(sqlite_path).suffix + ".pg-migration.json"),
|
||||
)
|
||||
)
|
||||
|
||||
if not sources:
|
||||
logger.info("Checkpoint dual-read enabled but no SQLite sources were opened; using Postgres only")
|
||||
yield saver
|
||||
return
|
||||
|
||||
proxy = DualReadCheckpointer(
|
||||
saver,
|
||||
sources,
|
||||
)
|
||||
try:
|
||||
yield proxy
|
||||
finally:
|
||||
await proxy.aclose_sources()
|
||||
return
|
||||
|
||||
raise ValueError(f"Unknown checkpointer type: {config.type!r}")
|
||||
|
||||
|
||||
+276
@@ -0,0 +1,276 @@
|
||||
"""SQLite-to-Postgres checkpoint migration and read fallback helpers.
|
||||
|
||||
The live gateway can wrap a Postgres checkpointer with ``DualReadCheckpointer``.
|
||||
New writes go to Postgres, reads prefer Postgres and fall back to legacy SQLite.
|
||||
Manual migration scripts reuse the copy helpers in this module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import sqlite3
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SQLiteCheckpointSource:
|
||||
path: str
|
||||
saver: Any
|
||||
context: Any
|
||||
progress_path: Path
|
||||
closed: bool = False
|
||||
|
||||
async def close(self) -> None:
|
||||
if self.closed:
|
||||
return
|
||||
self.closed = True
|
||||
try:
|
||||
await self.context.__aexit__(None, None, None)
|
||||
except Exception:
|
||||
logger.warning("Failed to close legacy SQLite checkpointer: %s", self.path, exc_info=True)
|
||||
|
||||
|
||||
def _sqlite_db_exists(path: str) -> bool:
|
||||
return Path(path).exists()
|
||||
|
||||
|
||||
def _sqlite_threads(path: str) -> list[str]:
|
||||
with sqlite3.connect(path) as conn:
|
||||
rows = conn.execute("SELECT DISTINCT thread_id FROM checkpoints ORDER BY thread_id").fetchall()
|
||||
return [str(row[0]) for row in rows if row and row[0]]
|
||||
|
||||
|
||||
def _sqlite_has_table(path: str, table: str) -> bool:
|
||||
with sqlite3.connect(path) as conn:
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?",
|
||||
(table,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
|
||||
def sqlite_recent_threads(path: str, *, since: str | None = None, until: str | None = None) -> list[str]:
|
||||
"""Return threads updated since an ISO timestamp when ``threads_meta`` exists."""
|
||||
|
||||
if not _sqlite_has_table(path, "threads_meta"):
|
||||
return []
|
||||
query = "SELECT thread_id FROM threads_meta"
|
||||
where: list[str] = []
|
||||
params: list[str] = []
|
||||
if since:
|
||||
where.append("updated_at >= ?")
|
||||
params.append(since)
|
||||
if until:
|
||||
where.append("updated_at <= ?")
|
||||
params.append(until)
|
||||
if where:
|
||||
query += " WHERE " + " AND ".join(where)
|
||||
query += " ORDER BY updated_at DESC"
|
||||
with sqlite3.connect(path) as conn:
|
||||
rows = conn.execute(query, tuple(params)).fetchall()
|
||||
return [str(row[0]) for row in rows if row and row[0]]
|
||||
|
||||
|
||||
def _sqlite_thread_checkpoint_count(path: str, thread_id: str) -> int:
|
||||
with sqlite3.connect(path) as conn:
|
||||
row = conn.execute(
|
||||
"SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?",
|
||||
(thread_id,),
|
||||
).fetchone()
|
||||
return int(row[0] if row else 0)
|
||||
|
||||
|
||||
def _load_progress(path: Path) -> set[str]:
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return set()
|
||||
except Exception:
|
||||
logger.warning("Ignoring unreadable checkpoint migration progress file: %s", path, exc_info=True)
|
||||
return set()
|
||||
migrated = payload.get("migrated_threads", [])
|
||||
return {str(x) for x in migrated if x}
|
||||
|
||||
|
||||
def _save_progress(path: Path, migrated: set[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||||
tmp.write_text(
|
||||
json.dumps({"migrated_threads": sorted(migrated)}, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
tmp.replace(path)
|
||||
|
||||
|
||||
async def _copy_thread_state(source, target, thread_id: str) -> bool:
|
||||
"""Copy one thread through the public LangGraph checkpointer API."""
|
||||
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
tuples = [t async for t in source.alist(config)]
|
||||
if not tuples:
|
||||
return False
|
||||
|
||||
for tup in reversed(tuples):
|
||||
configurable = (tup.config or {}).get("configurable", {}) or {}
|
||||
checkpoint_ns = configurable.get("checkpoint_ns", "") or ""
|
||||
parent_id = ((tup.parent_config or {}).get("configurable", {}) or {}).get("checkpoint_id")
|
||||
|
||||
put_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}}
|
||||
if parent_id:
|
||||
put_config["configurable"]["checkpoint_id"] = parent_id
|
||||
|
||||
new_versions = (tup.checkpoint or {}).get("channel_versions", {}) or {}
|
||||
await target.aput(put_config, tup.checkpoint, tup.metadata or {}, new_versions)
|
||||
|
||||
if tup.pending_writes:
|
||||
by_task: dict[str, list] = {}
|
||||
for task_id, channel, value in tup.pending_writes:
|
||||
by_task.setdefault(task_id, []).append((channel, value))
|
||||
checkpoint_id = (tup.checkpoint or {}).get("id")
|
||||
for task_id, writes in by_task.items():
|
||||
writes_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
}
|
||||
await target.aput_writes(writes_config, writes, task_id)
|
||||
return True
|
||||
|
||||
|
||||
async def copy_thread_state(source, target, thread_id: str) -> bool:
|
||||
return await _copy_thread_state(source, target, thread_id)
|
||||
|
||||
|
||||
async def _target_thread_checkpoint_count(target, thread_id: str) -> int:
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
count = 0
|
||||
async for _ in target.alist(config):
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
async def target_has_thread(target, thread_id: str) -> bool:
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
async for _ in target.alist(config, limit=1):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class DualReadCheckpointer:
|
||||
"""Postgres-first checkpointer with SQLite read fallback."""
|
||||
|
||||
def __init__(self, primary, sources: list[SQLiteCheckpointSource], *, sleep_seconds: float = 0.05):
|
||||
self.primary = primary
|
||||
self.sources = sources
|
||||
self.sleep_seconds = max(0.0, sleep_seconds)
|
||||
self._closed = False
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self.primary, name)
|
||||
|
||||
async def aclose_sources(self) -> None:
|
||||
self._closed = True
|
||||
for source in self.sources:
|
||||
await source.close()
|
||||
|
||||
async def aget_tuple(self, config):
|
||||
found = await self.primary.aget_tuple(config)
|
||||
if found is not None:
|
||||
return found
|
||||
for source in self.sources:
|
||||
if source.closed:
|
||||
continue
|
||||
found = await source.saver.aget_tuple(config)
|
||||
if found is not None:
|
||||
return found
|
||||
return None
|
||||
|
||||
async def aget(self, config):
|
||||
found = await self.primary.aget(config)
|
||||
if found is not None:
|
||||
return found
|
||||
for source in self.sources:
|
||||
if source.closed:
|
||||
continue
|
||||
found = await source.saver.aget(config)
|
||||
if found is not None:
|
||||
return found
|
||||
return None
|
||||
|
||||
async def alist(self, config, *, filter=None, before=None, limit=None):
|
||||
seen: set[tuple[str, str, str]] = set()
|
||||
yielded = 0
|
||||
|
||||
async for tup in self.primary.alist(config, filter=filter, before=before, limit=limit):
|
||||
key = self._tuple_key(tup)
|
||||
seen.add(key)
|
||||
yield tup
|
||||
yielded += 1
|
||||
if limit is not None and yielded >= limit:
|
||||
return
|
||||
|
||||
remaining = None if limit is None else max(0, limit - yielded)
|
||||
for source in self.sources:
|
||||
if source.closed:
|
||||
continue
|
||||
async for tup in source.saver.alist(config, filter=filter, before=before, limit=remaining):
|
||||
key = self._tuple_key(tup)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
yield tup
|
||||
yielded += 1
|
||||
if limit is not None:
|
||||
remaining = max(0, limit - yielded)
|
||||
if remaining <= 0:
|
||||
return
|
||||
|
||||
async def aput(self, *args, **kwargs):
|
||||
return await self.primary.aput(*args, **kwargs)
|
||||
|
||||
async def aput_writes(self, *args, **kwargs):
|
||||
return await self.primary.aput_writes(*args, **kwargs)
|
||||
|
||||
async def adelete_thread(self, thread_id: str) -> None:
|
||||
delete = getattr(self.primary, "adelete_thread", None)
|
||||
if delete is not None:
|
||||
await delete(thread_id)
|
||||
for source in self.sources:
|
||||
if source.closed:
|
||||
continue
|
||||
delete = getattr(source.saver, "adelete_thread", None)
|
||||
if delete is not None:
|
||||
await delete(thread_id)
|
||||
|
||||
@staticmethod
|
||||
def _tuple_key(tup) -> tuple[str, str, str]:
|
||||
configurable = (tup.config or {}).get("configurable", {}) or {}
|
||||
return (
|
||||
str(configurable.get("thread_id", "")),
|
||||
str(configurable.get("checkpoint_ns", "") or ""),
|
||||
str(configurable.get("checkpoint_id", "") or (tup.checkpoint or {}).get("id", "")),
|
||||
)
|
||||
|
||||
async def _verify_source(self, source: SQLiteCheckpointSource, threads: list[str]) -> bool:
|
||||
for thread_id in threads:
|
||||
source_count = await asyncio.to_thread(_sqlite_thread_checkpoint_count, source.path, thread_id)
|
||||
target_count = await _target_thread_checkpoint_count(self.primary, thread_id)
|
||||
if target_count < source_count:
|
||||
logger.warning(
|
||||
"Checkpoint migration verification failed: source=%s thread_id=%s source_count=%d target_count=%d",
|
||||
source.path,
|
||||
thread_id,
|
||||
source_count,
|
||||
target_count,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
@@ -187,8 +187,12 @@ checkpointer:
|
||||
# Docker 部署:把容器内 .deer-flow 目录挂卷出来即可持久化(撑住容器重启)。
|
||||
# 这里同时被「主 agent」和「AI 写作 graph」共用 —— LangGraph 用
|
||||
# (thread_id, checkpoint_ns) 做隔离,session_id 是 uuid,不会撞。
|
||||
type: sqlite
|
||||
connection_string: .deer-flow/data/checkpoints.db
|
||||
type: postgres
|
||||
connection_string: postgresql://postgres:postgres123@47.88.25.99:8432/deerflow
|
||||
dual_read_enabled: true
|
||||
dual_read_sqlite_paths:
|
||||
- backend/.deer-flow/data/deerflow.db
|
||||
- data/.deer-flow/data/checkpoints.db
|
||||
# AI 写作子系统的会话清理配置。
|
||||
# 业务表(ai_writing_sessions @ MySQL) + checkpointer(checkpoints.db @ SQLite) 联动清理。
|
||||
# 留意:清理走 cron 表达式调度,默认每天凌晨 3 点跑一次,按 updated_at 一刷拉清。
|
||||
|
||||
Reference in new issue
Block a user