1import math
2from collections import Counter
3
4documents = {
5    "runbook": "Restart the worker after rotating the OpenAI key.",
6    "pricing": "Track token usage by model and request id.",
7    "retrieval": "Hybrid search combines BM25 with vector similarity.",
8}
9
10def tokenize(text: str) -> list[str]:
11    return [term.strip(".,").lower() for term in text.split()]
12
13def score(query: str, text: str) -> float:
14    query_terms = Counter(tokenize(query))
15    doc_terms = Counter(tokenize(text))
16    shared = set(query_terms) & set(doc_terms)
17    numerator = sum(query_terms[term] * doc_terms[term] for term in shared)
18    query_norm = math.sqrt(sum(value * value for value in query_terms.values()))
19    doc_norm = math.sqrt(sum(value * value for value in doc_terms.values()))
20    return numerator / (query_norm * doc_norm or 1)
21
22def retrieve(query: str, k: int = 2) -> list[tuple[str, float]]:
23    ranked = [(doc_id, round(score(query, text), 3)) for doc_id, text in documents.items()]
24    return sorted(ranked, key=lambda row: row[1], reverse=True)[:k]
25
26print(retrieve("hybrid search token usage"))