from haystack import Document
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder

embedder = SentenceTransformersTextEmbedder(
    model="sentence-transformers/all-MiniLM-L6-v2",
    normalize_embeddings=True,
    progress_bar=False,
)

embedder.warm_up()
query_result = embedder.run(text="password reset approval policy")
query_embedding = query_result["embedding"]

document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
document_store.write_documents(
    [
        Document(
            content="Password reset requests require help desk approval.",
            meta={"name": "password-reset-policy"},
            embedding=query_embedding,
        )
    ]
)

retriever = InMemoryEmbeddingRetriever(document_store=document_store, top_k=1)
match = retriever.run(query_embedding=query_embedding)["documents"][0]

print(f"embedding_dimensions={len(query_embedding)}")
print(
    f"first_values={query_embedding[0]:.4f},"
    f"{query_embedding[1]:.4f},"
    f"{query_embedding[2]:.4f}"
)
print(f"normalized_l2={sum(value * value for value in query_embedding):.6f}")
print(f"retriever_top_document={match.meta['name']}")