rag/retriever.py

from irpr.deps import search

def retrieve(query: str, top_k=8):
    return search(query, top_k=top_k)

def format_citations(hits):
    url_to_id = {}
    uniq = []
    for h in hits:
        u = h["source_url"]
        if u not in url_to_id:
            url_to_id[u] = len(url_to_id) + 1
            uniq.append((url_to_id[u], u))
    return uniq