from pathlib import Path
import shutil

from llama_index.core import (
    Document,
    MockEmbedding,
    Settings,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)


Settings.embed_model = MockEmbedding(embed_dim=8)
storage_dir = Path("storage")
shutil.rmtree(storage_dir, ignore_errors=True)


original = Document(
    text="Escalate payment failures to the billing queue.",
    id_="support-policy",
)
index = VectorStoreIndex.from_documents([original])
index.storage_context.persist(persist_dir=storage_dir)


updated = Document(
    text="Escalate payment failures to the revenue operations queue.",
    id_="support-policy",
)
storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
index = load_index_from_storage(storage_context)
refresh_flags = index.refresh_ref_docs([updated])
index.storage_context.persist(persist_dir=storage_dir)


storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
index = load_index_from_storage(storage_context)
nodes = index.as_retriever(similarity_top_k=1).retrieve("payment failures")

print(f"Refresh flags: {refresh_flags}")
print(f"Reloaded text: {nodes[0].node.get_content(metadata_mode='none')}")