ResearchMind / src /persistence.py
Vivek6041's picture
Upload 7 files
86a4642 verified
Raw History Blame Contribute Delete
1.72 kB
import json
from pathlib import Path
from typing import Optional
from src.config import VECTORSTORE_DIR, PAPERS_DIR
def get_vectorstore_path(paper_id: str) -> Path:
"""Get the path where a vectorstore for a paper is saved."""
return VECTORSTORE_DIR / f"{paper_id}.faiss"
def get_metadata_path(paper_id: str) -> Path:
"""Get the path where paper metadata is saved."""
return VECTORSTORE_DIR / f"{paper_id}_metadata.json"
def save_paper_metadata(paper_id: str, filename: str, chunks_count: int):
"""Save paper metadata (filename, chunk count, etc.)."""
metadata = {
"paper_id": paper_id,
"filename": filename,
"chunks_count": chunks_count,
}
with open(get_metadata_path(paper_id), "w") as f:
json.dump(metadata, f, indent=2)
def load_paper_metadata(paper_id: str) -> Optional[dict]:
"""Load paper metadata if it exists."""
path = get_metadata_path(paper_id)
if path.exists():
with open(path, "r") as f:
return json.load(f)
return None
def list_saved_papers() -> list[dict]:
"""List all saved papers with their metadata."""
papers = []
for metadata_file in VECTORSTORE_DIR.glob("*_metadata.json"):
with open(metadata_file, "r") as f:
papers.append(json.load(f))
return papers
def delete_paper(paper_id: str):
"""Delete a paper's vectorstore and metadata."""
vectorstore_path = get_vectorstore_path(paper_id)
metadata_path = get_metadata_path(paper_id)
if vectorstore_path.exists():
import shutil
shutil.rmtree(vectorstore_path)
if metadata_path.exists():
metadata_path.unlink()