Query rewriting, HyDE and query routing

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
 
chunks = [
    "The Model X battery lasts about 10 hours of video playback.",
    "The Model X screen is a 14-inch OLED panel.",
    "Battery replacements cost 89 pounds and take 3 days.",
]
vec = TfidfVectorizer().fit(chunks)
def search(q):
    s = cosine_similarity(vec.transform([q]), vec.transform(chunks))[0]
    return round(float(s.max()), 2), chunks[s.argmax()]
 
history = "User: I'm looking at the Model X laptop."
follow_up = "how long does it last?"
 
# 1) Query rewriting: turn a vague follow-up into a standalone search query
#    (an LLM does this in real life; written by hand here).
rewritten = "How long does the Model X laptop battery last?"
print("raw     :", search(follow_up))
print("rewrite :", search(rewritten))
 
# 2) HyDE: let a model write a hypothetical answer, then search with THAT text.
hypothetical = "The Model X battery lasts around 10 hours on a full charge."
print("HyDE    :", search(hypothetical))
 
# 3) Routing: decide which index (or tool) should answer before searching.
def route(q):
    q = q.lower()
    return "orders_db" if "order" in q else "pricing_table" if "cost" in q or "price" in q else "product_docs"
for q in ["Where is my order 1234?", "What does a new battery cost?", "Is the screen OLED?"]:
    print(f"route   : {q!r} -> {route(q)}")

raw     : (0.0, 'The Model X battery lasts about 10 hours of video playback.')
rewrite : (0.45, 'The Model X battery lasts about 10 hours of video playback.')
HyDE    : (0.74, 'The Model X battery lasts about 10 hours of video playback.')
route   : 'Where is my order 1234?' -> orders_db
route   : 'What does a new battery cost?' -> pricing_table
route   : 'Is the screen OLED?' -> product_docs