import json

import numpy as np

from ingest import embed

CHUNKS = [json.loads(line) for line in open("chunks.jsonl")]
VECTORS = np.load("embeddings.npy")


def retrieve(query, k=20):
    q = embed([query], "search_query: ")[0]
    scores = VECTORS @ q
    top = np.argsort(-scores)[:k]
    return [dict(CHUNKS[i], similarity=float(scores[i])) for i in top]


def retrieve_many(queries, k=20):
    best = {}
    for query in queries:
        for chunk in retrieve(query, k if len(queries) == 1 else 15):
            if chunk["id"] not in best or chunk["similarity"] > best[chunk["id"]]["similarity"]:
                best[chunk["id"]] = chunk
    return sorted(best.values(), key=lambda c: -c["similarity"])