File size: 3,357 Bytes
20f1ed0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """
utils/file_registry.py
======================
Provides SHA-256 based duplicate-PDF detection for the ingestion pipeline.
A JSON registry is persisted to disk so that already-indexed documents are
remembered across application restarts. Only the ingestion path reads or
writes this module; retrieval, reranking, and the LLM are unaffected.
"""
import hashlib
import json
import os
from typing import Dict, Any
# Default location for the on-disk registry
DEFAULT_REGISTRY_PATH = "./data/processed_files.json"
# ---------------------------------------------------------------------------
# Hashing
# ---------------------------------------------------------------------------
def compute_file_hash(file_path: str) -> str:
"""
Return the SHA-256 hex digest of the file at *file_path*.
Reads in 8 KB chunks so large PDFs do not exhaust memory.
"""
sha256 = hashlib.sha256()
with open(file_path, "rb") as f:
while chunk := f.read(8192):
sha256.update(chunk)
return sha256.hexdigest()
# ---------------------------------------------------------------------------
# Registry I/O
# ---------------------------------------------------------------------------
def load_registry(registry_path: str = DEFAULT_REGISTRY_PATH) -> Dict[str, Any]:
"""
Load the processed-files registry from *registry_path*.
Returns an empty dict if the file does not yet exist.
Registry schema::
{
"<sha256_hex>": {
"filename": "<original filename>",
"num_chunks": <int>
},
...
}
"""
if os.path.exists(registry_path):
with open(registry_path, "r", encoding="utf-8") as f:
return json.load(f)
return {}
def save_registry(registry: Dict[str, Any], registry_path: str = DEFAULT_REGISTRY_PATH) -> None:
"""Persist *registry* to *registry_path* as pretty-printed JSON."""
os.makedirs(os.path.dirname(registry_path), exist_ok=True)
with open(registry_path, "w", encoding="utf-8") as f:
json.dump(registry, f, indent=2)
def register_file(
registry: Dict[str, Any],
file_hash: str,
filename: str,
num_chunks: int,
registry_path: str = DEFAULT_REGISTRY_PATH,
) -> None:
"""
Add *file_hash* → metadata to *registry* and immediately persist to disk.
Args:
registry: The in-memory registry dict (mutated in-place).
file_hash: SHA-256 hex digest of the file.
filename: Human-readable filename for logging / UI feedback.
num_chunks: Number of chunks produced during chunking.
registry_path: Where to persist the JSON registry.
"""
registry[file_hash] = {
"filename": filename,
"num_chunks": num_chunks,
}
save_registry(registry, registry_path)
def is_duplicate(registry: Dict[str, Any], file_hash: str) -> bool:
"""Return *True* if *file_hash* is already present in *registry*."""
return file_hash in registry
def clear_registry(registry_path: str = DEFAULT_REGISTRY_PATH) -> None:
"""
Delete the on-disk registry file and reset state.
Called when the user clears the FAISS index so that previously-indexed
hashes are no longer treated as duplicates.
"""
if os.path.exists(registry_path):
os.remove(registry_path)
|