from sentence_transformers import SentenceTransformer
from sentence_transformers import util

model = SentenceTransformer(
    "all-MiniLM-L6-v2"
)

def flatten_knowledge(data):

    texts = []

    for section in data.values():

        if isinstance(section, list):

            texts.extend(section)

    return texts

def semantic_search(query, knowledge):

    texts = flatten_knowledge(knowledge)

    corpus_embeddings = model.encode(
        texts,
        convert_to_tensor=True
    )

    query_embedding = model.encode(
        query,
        convert_to_tensor=True
    )

    hits = util.semantic_search(
        query_embedding,
        corpus_embeddings,
        top_k=3
    )[0]

    answer = []

    for hit in hits:

        answer.append(
            texts[hit["corpus_id"]]
        )

    return "\n\n".join(answer)