Coding Round

Example

Plain Python, no numpy. Sort by score descending and break ties by ID so the result is stable.

import math

def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na, nb = math.sqrt(sum(x * x for x in a)), math.sqrt(sum(y * y for y in b))
    return 0.0 if na == 0 or nb == 0 else dot / (na * nb)

def top_k(query_vec, index, k=2):
    scored = [(round(cosine(query_vec, vec), 3), doc_id) for doc_id, vec in index.items()]
    scored.sort(key=lambda t: (-t[0], t[1]))
    return scored[:k]

index = {"returns": [0.9, 0.1, 0.0], "refunds": [0.7, 0.6, 0.1],
         "shipping": [0.0, 0.2, 0.9], "warranty": [0.5, 0.0, 0.4]}
print(top_k([1.0, 0.3, 0.0], index, k=3))
print(top_k([0.0, 0.0, 0.0], index, k=1))   # zero vector edge case

Output

[(0.984, 'returns'), (0.909, 'refunds'), (0.748, 'warranty')]
[(0.0, 'refunds')]

Solution

def chunk_words(text, size, overlap):
    if not 0 <= overlap < size:
        raise ValueError("need 0 <= overlap < size")
    words = text.split()
    step = size - overlap
    chunks = []
    for start in range(0, len(words), step):
        chunks.append(" ".join(words[start:start + size]))
        if start + size >= len(words):
            break
    return chunks

text = ("Most items can be returned within 10 days of delivery. Electronics have a "
        "7 day replacement only policy. Refunds reach the original payment method in 5-7 working days.")
for i, c in enumerate(chunk_words(text, size=10, overlap=3)):
    print(i, c)

Output

0 Most items can be returned within 10 days of delivery.
1 days of delivery. Electronics have a 7 day replacement only
2 day replacement only policy. Refunds reach the original payment method
3 original payment method in 5-7 working days.

Solution

import math
from collections import Counter

def bm25_scores(query, docs, k1=1.5, b=0.75):
    toks = [d.lower().split() for d in docs]
    n, avgdl = len(toks), sum(len(t) for t in toks) / len(toks)
    df = Counter(w for t in toks for w in set(t))
    scores = []
    for t in toks:
        tf, s = Counter(t), 0.0
        for w in query.lower().split():
            if w not in tf:
                continue
            idf = math.log((n - df[w] + 0.5) / (df[w] + 0.5) + 1)
            s += idf * tf[w] * (k1 + 1) / (tf[w] + k1 * (1 - b + b * len(t) / avgdl))
        scores.append(round(s, 3))
    return scores

docs = ["cod available up to 50000 rupees",
        "refund to original payment method in 5 to 7 days",
        "cod orders get refund to bank account"]
print(bm25_scores("cod refund", docs))

Output

[0.521, 0.413, 0.978]

Example

Token counts are estimated as about 4 characters per token here; in production use the provider's token counting endpoint.

def est_tokens(text):
    return max(1, len(text) // 4)

def trim_history(system, messages, budget):
    used = est_tokens(system)
    kept = []
    for msg in reversed(messages):                 # newest first
        cost = est_tokens(msg["content"])
        if used + cost > budget:
            break
        kept.append(msg)
        used += cost
    kept.reverse()
    while kept and kept[0]["role"] != "user":      # history must start with a user turn
        kept.pop(0)
    return kept, used

system = "You are ShopEase support. Be brief."
messages = [
    {"role": "user", "content": "Hi, I ordered a pressure cooker last week, order SE1077."},
    {"role": "assistant", "content": "Thanks! SE1077 is out for delivery and should arrive today."},
    {"role": "user", "content": "Also, what is the return window for shoes?"},
    {"role": "assistant", "content": "Shoes can be returned within 10 days with tags attached."},
    {"role": "user", "content": "And for earbuds?"},
]
kept, used = trim_history(system, messages, budget=45)
print("kept", len(kept), "of", len(messages), "messages, ~", used, "tokens")
for m in kept:
    print(f"  {m['role']}: {m['content']}")

Output

kept 3 of 5 messages, ~ 36 tokens
  user: Also, what is the return window for shoes?
  assistant: Shoes can be returned within 10 days with tags attached.
  user: And for earbuds?

Example

fake_model is a stand-in: the first reply has a code fence and a trailing comma, the second misses a field, the third is valid.

import json, re

REPLIES = iter([
    'Sure! Here it is:\n~~~json\n{"intent": "refund_status", "order_id": "SE1042",}\n~~~',
    '{"intent": "refund_status"}',
    '{"intent": "refund_status", "order_id": "SE1042", "urgent": false}',
])
def fake_model(prompt):  # stand-in for the model
    return next(REPLIES)

REQUIRED = {"intent": str, "order_id": str, "urgent": bool}

def repair(text):
    m = re.search(r"\{.*\}", text, re.S)           # keep only the outermost object
    if not m:
        raise ValueError("no JSON object found")
    return re.sub(r",\s*([}\]])", r"\1", m.group())  # remove trailing commas

def validate(obj):
    for key, typ in REQUIRED.items():
        if not isinstance(obj.get(key), typ):
            raise ValueError(f"field '{key}' missing or not {typ.__name__}")
    return obj

def extract(prompt, max_attempts=3):
    for attempt in range(1, max_attempts + 1):
        reply = fake_model(prompt)
        try:
            result = validate(json.loads(repair(reply)))
            print(f"attempt {attempt}: ok")
            return result
        except (ValueError, json.JSONDecodeError) as e:
            print(f"attempt {attempt}: {e}")
            prompt += f"\nYour last reply was invalid: {e}. Return only valid JSON."
    raise RuntimeError("could not get valid JSON")

print(extract("Classify this ticket as JSON: 'SE1042 refund?'"))

Output

attempt 1: field 'urgent' missing or not bool
attempt 2: field 'order_id' missing or not str
attempt 3: ok
{'intent': 'refund_status', 'order_id': 'SE1042', 'urgent': False}

Example

scripted_model is a stand-in for the API: it returns blocks shaped like Claude's tool_use / text content.

def get_order(order_id):
    return {"SE1042": "delivered on 12 Oct"}.get(order_id, "not found")

def days_since(date_text):
    return "9 days"           # fixed for the demo

TOOLS = {"get_order": get_order, "days_since": days_since}

def scripted_model(messages):  # stand-in for client.messages.create
    n_results = sum(1 for m in messages if m["role"] == "user" and isinstance(m["content"], list))
    if n_results == 0:
        return "tool_use", [{"type": "tool_use", "id": "t1", "name": "get_order", "input": {"order_id": "SE1042"}}]
    if n_results == 1:
        return "tool_use", [{"type": "tool_use", "id": "t2", "name": "days_since", "input": {"date_text": "12 Oct"}}]
    return "end_turn", [{"type": "text", "text": "SE1042 was delivered 9 days ago, so you still have 1 day to return it."}]

def run_agent(question, max_steps=5):
    messages = [{"role": "user", "content": question}]
    for step in range(max_steps):
        stop, content = scripted_model(messages)
        messages.append({"role": "assistant", "content": content})
        if stop == "end_turn":
            return next(b["text"] for b in content if b["type"] == "text")
        results = []
        for b in content:
            if b["type"] == "tool_use":
                fn = TOOLS.get(b["name"])
                out = fn(**b["input"]) if fn else f"unknown tool {b['name']}"
                print(f"step {step + 1}: {b['name']}({b['input']}) -> {out}")
                results.append({"type": "tool_result", "tool_use_id": b["id"], "content": out,
                                "is_error": fn is None})
        messages.append({"role": "user", "content": results})
    return "Stopped: step limit reached."

print(run_agent("Can I still return SE1042?"))

Output

step 1: get_order({'order_id': 'SE1042'}) -> delivered on 12 Oct
step 2: days_since({'date_text': '12 Oct'}) -> 9 days
SE1042 was delivered 9 days ago, so you still have 1 day to return it.

Example

Bag-of-words vectors stand in for an embedding model.

import math, re
from collections import Counter

def embed(text):  # stand-in for an embedding model
    stop = {"what", "is", "the", "a", "for", "how", "do", "i", "my", "to", "can"}
    return Counter(w for w in re.findall(r"[a-z]+", text.lower()) if w not in stop)

def cosine(a, b):
    dot = sum(a[w] * b[w] for w in a)
    norm = math.sqrt(sum(v * v for v in a.values())) * math.sqrt(sum(v * v for v in b.values()))
    return dot / norm if norm else 0.0

class SemanticCache:
    def __init__(self, threshold):
        self.threshold, self.entries = threshold, []

    def lookup(self, question):
        q = embed(question)
        scored = [(cosine(q, vec), ans) for vec, ans in self.entries]
        best = max(scored, default=(0.0, None))
        return (best[1] if best[0] >= self.threshold else None), round(best[0], 3)

    def store(self, question, answer):
        self.entries.append((embed(question), answer))

calls = 0
def answer(question, cache):
    global calls
    hit, sim = cache.lookup(question)
    if hit:
        return f"cache ({sim}) -> {hit}"
    calls += 1
    reply = f"model answer #{calls}"           # stand-in for the LLM call
    cache.store(question, reply)
    return f"model ({sim}) -> {reply}"

cache = SemanticCache(threshold=0.75)
for q in ["What is the return window for shoes?", "return window for shoes",
          "What is the refund time for UPI?", "UPI refund time"]:
    print(f"{q:38} {answer(q, cache)}")
print("model calls:", calls)

Output

What is the return window for shoes?   model (0.0) -> model answer #1
return window for shoes                cache (1.0) -> model answer #1
What is the refund time for UPI?       model (0.0) -> model answer #2
UPI refund time                        cache (1.0) -> model answer #2
model calls: 2

Solution

import math

def softmax(logits, temperature=1.0):
    if temperature <= 0:
        raise ValueError("temperature must be > 0 (use argmax for greedy decoding)")
    scaled = [x / temperature for x in logits]
    m = max(scaled)                                   # subtract max for numerical stability
    exps = [math.exp(x - m) for x in scaled]
    total = sum(exps)
    return [e / total for e in exps]

tokens = ["delivered", "shipped", "lost", "banana"]
logits = [3.0, 2.0, 0.5, -1.0]
for t in (0.2, 1.0, 2.0):
    probs = softmax(logits, t)
    print(f"T={t}: " + "  ".join(f"{tok}={p:.3f}" for tok, p in zip(tokens, probs)))
print("overflow-safe:", [round(p, 3) for p in softmax([1000.0, 999.0])])

Output

T=0.2: delivered=0.993  shipped=0.007  lost=0.000  banana=0.000
T=1.0: delivered=0.681  shipped=0.251  lost=0.056  banana=0.012
T=2.0: delivered=0.493  shipped=0.299  lost=0.141  banana=0.067
overflow-safe: [0.731, 0.269]

Solution

def precision_at_k(ranked, relevant, k):
    return sum(1 for d in ranked[:k] if d in relevant) / k

def recall_at_k(ranked, relevant, k):
    return sum(1 for d in ranked[:k] if d in relevant) / len(relevant) if relevant else 0.0

def mrr(runs):
    total = 0.0
    for ranked, relevant in runs:
        total += next((1 / i for i, d in enumerate(ranked, 1) if d in relevant), 0.0)
    return total / len(runs)

runs = [
    (["d3", "d1", "d7"], {"d1"}),
    (["d2", "d5", "d9"], {"d2", "d9"}),
    (["d4", "d6", "d8"], {"d1"}),
]
for ranked, rel in runs:
    print(ranked, sorted(rel), f"P@2={precision_at_k(ranked, rel, 2):.2f}", f"R@2={recall_at_k(ranked, rel, 2):.2f}")
print(f"MRR={mrr(runs):.3f}")

Output

['d3', 'd1', 'd7'] ['d1'] P@2=0.50 R@2=1.00
['d2', 'd5', 'd9'] ['d2', 'd9'] P@2=0.50 R@2=0.50
['d4', 'd6', 'd8'] ['d1'] P@2=0.00 R@2=0.00
MRR=0.500