Joeven — Learn Autonomous AI Agents

Python loads on first run

QUERY = "how long do INV-17 refunds take"
POOL = [
    "The cafeteria menu is pizza on Fridays.",
    "Refunds for INV-* take 5-7 days. Never refund in cash.",
    "INV-17 is an open invoice for 40 dollars.",
    "Raise worker memory when the runner hits OOM.",
    "Refunds: 5-7 days. INV-17 included. Check status in billing.md.",
]

def retrieve_stage1(q, docs, k=4):
    qw = set(q.lower().split())
    scored = [(len(qw & set(d.lower().split())), d) for d in docs]
    scored.sort(reverse=True)
    return [d for _, d in scored[:k]]

def rerank_score(q, d):
    qw = q.lower().split()
    dl = d.lower()
    overlap = sum(1 for w in qw if w in dl)
    if "5-7" in dl or "days" in dl:
        overlap += 2
    if "inv-17" in dl and "refund" in q.lower():
        overlap += 2
    return overlap

def rerank(q, docs, n=2):
    ranked = sorted(((rerank_score(q, d), d) for d in docs), reverse=True)
    return ranked[:n]

pile = retrieve_stage1(QUERY, POOL, k=4)
print("stage1:")
for d in pile:
    print(" ", d[:70])
print("stage2:")
for score, d in rerank(QUERY, pile, n=2):
    print(" ", score, d[:70])

Run to execute this in your browser. Nothing is sent to a server.