mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat: add persistent vector build state helpers
This commit is contained in:
parent
28562811b1
commit
64ee50a080
2 changed files with 93 additions and 0 deletions
|
|
@ -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)
|
||||
|
|
|
|||
41
tests/unit/memory/test_vector_db.py
Normal file
41
tests/unit/memory/test_vector_db.py
Normal 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"
|
||||
Loading…
Reference in a new issue