from sentence_transformers import SentenceTransformer
from sentence_transformers.util import semantic_search


model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

chunks = [
    {
        "id": "kb-001",
        "title": "Reset a forgotten password",
        "text": "Open Account settings, send a reset email, and choose a new password.",
    },
    {
        "id": "kb-002",
        "title": "Download invoice receipts",
        "text": "Billing admins can download paid invoice receipts from billing history.",
    },
    {
        "id": "kb-003",
        "title": "Rotate API tokens",
        "text": "Create a replacement token, update the integration, and revoke the old token.",
    },
]

query = "How can a user reset a forgotten password?"
documents = [f"{chunk['title']}: {chunk['text']}" for chunk in chunks]

document_embeddings = model.encode_document(
    documents,
    normalize_embeddings=True,
    convert_to_tensor=True,
    show_progress_bar=False,
)
query_embedding = model.encode_query(
    [query],
    normalize_embeddings=True,
    convert_to_tensor=True,
    show_progress_bar=False,
)

hits = semantic_search(query_embedding, document_embeddings, top_k=2)[0]
retrieved_chunks = [chunks[hit["corpus_id"]] for hit in hits]
context = "\n\n".join(
    f"[{chunk['id']}] {chunk['title']}\n{chunk['text']}"
    for chunk in retrieved_chunks
)

print(f"query: {query}")
print("retrieved chunks:")
for rank, hit in enumerate(hits, start=1):
    chunk = chunks[hit["corpus_id"]]
    print(f"{rank}. {chunk['id']} title={chunk['title']}")

print("\nrag context:")
print(context)

if retrieved_chunks[0]["id"] != "kb-001":
    raise SystemExit(f"unexpected top chunk: {retrieved_chunks[0]['id']}")
if "[kb-001] Reset a forgotten password" not in context:
    raise SystemExit("password-reset context is missing")