LlamaIndex StorageContext — index reloaded with mismatched embed_model

import json
from pathlib import Path
from llama_index.core import VectorStoreIndex, StorageContext, Settings, load_index_from_storage
from llama_index.embeddings.openai import OpenAIEmbedding

def persist_index(index, persist_dir: str):
    """Persist index AND its embedding config for safe reload."""
    p = Path(persist_dir)
    p.mkdir(parents=True, exist_ok=True)

    # Standard index persist
    index.storage_context.persist(persist_dir=persist_dir)

    # Also record the embed_model config
    embed_model = index._embed_model
    config = {
        "embed_model_class": type(embed_model).__name__,
        "embed_model_name": getattr(embed_model, "model_name", None) or getattr(embed_model, "model", None),
        "embed_dim": getattr(embed_model, "embed_dim", None),
        # Include any provider-specific config that affects embeddings
        "additional_kwargs": getattr(embed_model, "additional_kwargs", {}),
    }
    (p / "embed_config.json").write_text(json.dumps(config, indent=2))
    print(f"Persisted index and config to {persist_dir}")


def load_index_safely(persist_dir: str, embed_model=None):
    """Load index, verifying embedding config matches what was persisted."""
    p = Path(persist_dir)

    # Load the config
    config_path = p / "embed_config.json"
    if not config_path.exists():
        raise ValueError(
            f"No embed_config.json in {persist_dir}. This index was persisted "
            f"without safety metadata; cannot verify embed_model match."
        )
    saved_config = json.loads(config_path.read_text())

    # Verify current embed_model matches
    if embed_model is None:
        embed_model = Settings.embed_model

    current_class = type(embed_model).__name__
    current_model = getattr(embed_model, "model_name", None) or getattr(embed_model, "model", None)

    if current_class != saved_config["embed_model_class"]:
        raise ValueError(
            f"embed_model class mismatch: index was built with {saved_config['embed_model_class']}, "
            f"but current is {current_class}. Retrieval will be broken."
        )
    if current_model != saved_config["embed_model_name"]:
        raise ValueError(
            f"embed_model name mismatch: index was built with {saved_config['embed_model_name']}, "
            f"but current is {current_model}. Retrieval will be broken."
        )

    storage_context = StorageContext.from_defaults(persist_dir=persist_dir)
    return load_index_from_storage(storage_context, embed_model=embed_model)

# Usage
embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.embed_model = embed_model

# Persist
persist_index(index, "./storage/my_index")

# Later — safe load
index = load_index_safely("./storage/my_index", embed_model=embed_model)

from llama_index.core import StorageContext, load_index_from_storage, Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# --- Explicit reload with all config ---
def rebuild_query_engine(persist_dir: str):
    """Rebuild query engine from persisted storage with explicit config."""

    # Recreate the exact embed_model used at ingest
    embed_model = OpenAIEmbedding(
        model="text-embedding-3-small",
        embed_batch_size=32,
    )

    llm = OpenAI(model="gpt-4o-mini")

    # Set as Settings (helpful for downstream)
    Settings.embed_model = embed_model
    Settings.llm = llm

    # Load with explicit reference (do not rely on Settings)
    storage_context = StorageContext.from_defaults(persist_dir=persist_dir)
    index = load_index_from_storage(
        storage_context,
        embed_model=embed_model,      # ← explicit
        llm=llm,                       # ← explicit
    )

    return index.as_query_engine(similarity_top_k=6, llm=llm)

# --- For remote vector stores, reload the store connection ---
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

def reload_with_remote_vector_store(persist_dir: str):
    # Reconnect to the remote vector store
    client = chromadb.PersistentClient(path="./chroma_db")
    collection = client.get_collection("my_docs")
    vector_store = ChromaVectorStore(chroma_collection=collection)

    # persist_dir only contains docstore + index_store; vector_store is remote
    storage_context = StorageContext.from_defaults(
        vector_store=vector_store,     # ← from remote
        persist_dir=persist_dir,        # ← docstore, index_store from local
    )

    embed_model = OpenAIEmbedding(model="text-embedding-3-small")
    return load_index_from_storage(storage_context, embed_model=embed_model)

from llama_index.core import VectorStoreIndex, load_index_from_storage
from typing import Tuple

# --- Canonical smoke test defined at ingest time, stored with index ---

def run_smoke_test(index, canonical_query: str, expected_source: str,
                  min_score: float = 0.4) -> Tuple[bool, str]:
    """Return (passed, message). Raises on hard failures like dimension mismatch."""
    try:
        retriever = index.as_retriever(similarity_top_k=5)
        nodes = retriever.retrieve(canonical_query)
    except ValueError as e:
        if "dimension" in str(e).lower():
            return False, f"DIMENSION MISMATCH: {e}"
        raise

    if not nodes:
        return False, "Retriever returned 0 nodes — likely empty index or filter issue"

    top_score = nodes[0].score
    if top_score < min_score:
        return False, (
            f"Top score {top_score:.3f} below threshold {min_score:.3f} — "
            f"embed_model may be different from ingest"
        )

    top_source = nodes[0].metadata.get("source", "")
    if expected_source not in top_source:
        return False, (
            f"Top result {top_source} does not match expected {expected_source}. "
            f"Index may be corrupted or embed_model wrong."
        )

    return True, f"OK — top score {top_score:.3f}, source {top_source}"

# --- Usage after every load ---
index = load_index_safely("./storage/my_index", embed_model=embed_model)

passed, msg = run_smoke_test(
    index,
    canonical_query="canonical smoke test query",
    expected_source="doc_smoke_canonical_source",
    min_score=0.4,
)
if not passed:
    raise RuntimeError(f"Index smoke test failed: {msg}")
print(f"Smoke test: {msg}")