search.py

from utils import load_model, load_collection

search_params = search_params = {
    "metric_type": "COSINE",  # L2
    "params": {"level": 2}
}

def search_chunks(query, model_name, collection_name, topk=5):
    model = load_model(model_name)
    collection = load_collection(collection_name)
    embed_query = model.encode(query)
    results = collection.search(
        [embed_query],
        anns_field="chunk_embedding",
        param=search_params,
        limit=topk,
        guarantee_timestamp=1,
        output_fields=['chunk_ref', 'chunk_text'] #
    )
    return results[0]