feat: add persistent vector build state helpers

This commit is contained in:
Research Assistant 2026-05-15 01:13:12 +08:00
parent 28562811b1
commit 64ee50a080
2 changed files with 93 additions and 0 deletions

View file

@ -282,6 +282,58 @@ def retrieve_chunks(vault: Path, query: str, limit: int = 5, expand: bool = True
return chunks
# ── Persistent build state ─────────────────────────────────────────────────
def get_vector_build_state_path(vault: Path) -> Path:
"""Return path to vector-build-state.json."""
from paperforge.config import paperforge_paths
paths = paperforge_paths(vault)
index_dir = (paths.get("memory_db", paths.get("index", vault / "System" / "PaperForge"))).parent
return index_dir / "vector-build-state.json"
def read_vector_build_state(vault: Path) -> dict:
"""Read persisted build state. Returns default idle state if file missing."""
path = get_vector_build_state_path(vault)
if not path.exists():
return {
"status": "idle",
"current": 0,
"total": 0,
"paper_id": "",
"last_update": "",
"started_at": "",
"finished_at": "",
"resume_supported": True,
"mode": "local",
"model": "BAAI/bge-small-en-v1.5",
"message": "",
"pid": 0,
}
try:
return json.loads(path.read_text(encoding="utf-8"))
except Exception:
return {"status": "idle", "current": 0, "total": 0, "paper_id": ""}
def write_vector_build_state(vault: Path, state: dict) -> None:
"""Persist build state to disk atomically."""
path = get_vector_build_state_path(vault)
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(".tmp")
tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
tmp.replace(path)
def mark_vector_build_state(vault: Path, **fields) -> dict:
"""Update specific fields in the build state file. Returns updated state."""
state = read_vector_build_state(vault)
state.update(fields)
write_vector_build_state(vault, state)
return state
def get_embed_status(vault: Path) -> dict:
"""Get vector DB status."""
db_path = get_vector_db_path(vault)

View file

@ -0,0 +1,41 @@
from __future__ import annotations
from paperforge.memory.vector_db import (
get_vector_build_state_path,
read_vector_build_state,
write_vector_build_state,
mark_vector_build_state,
)
def test_vector_build_state_defaults_when_missing(tmp_path):
vault = tmp_path / "vault"
vault.mkdir()
state = read_vector_build_state(vault)
assert state["status"] == "idle"
assert state["current"] == 0
def test_vector_build_state_roundtrip(tmp_path):
vault = tmp_path / "vault"
vault.mkdir()
original = {
"status": "running", "current": 5, "total": 20,
"paper_id": "ABC12345", "last_update": "now",
"started_at": "", "finished_at": "",
"resume_supported": True, "mode": "local",
"model": "bge-small", "message": "", "pid": 0,
}
write_vector_build_state(vault, original)
loaded = read_vector_build_state(vault)
assert loaded == original
def test_mark_vector_build_state_updates_field(tmp_path):
vault = tmp_path / "vault"
vault.mkdir()
write_vector_build_state(vault, {"status": "idle", "current": 0, "total": 0})
mark_vector_build_state(vault, current=10, paper_id="XYZ")
state = read_vector_build_state(vault)
assert state["current"] == 10
assert state["paper_id"] == "XYZ"