Agentic RAG take-home project

"""Practice corpus: internal notes of a fictional robotics company, and
questions that need one or two lookups. Do not edit."""

DOCS = {
    "d1": "The Atlas scheduler assigns delivery jobs to robots. It is owned by the Fleet Platform team.",
    "d2": "The Fleet Platform team is led by Priya Raman and sits in the Austin office.",
    "d3": "Beacon is the telemetry pipeline. It is owned by the Observability team and stores data for 30 days.",
    "d4": "The Observability team is led by Tomas Ekberg and sits in the Lisbon office.",
    "d5": "Robots in the Austin office charge at docks built by Voltline. Each dock charges two robots at once.",
    "d6": "Robots in the Lisbon office charge at docks built by Ampere Works. Each dock charges four robots at once.",
    "d7": "The Austin office runs 120 robots. The Lisbon office runs 80 robots.",
    "d8": "Atlas retries a failed job three times before paging the on-call engineer.",
    "d9": "Incident 412: Beacon dropped data for six hours after a disk filled up. The fix was a retention alarm.",
    "d10": "The cafeteria in the Austin office serves tacos on Fridays.",
}

# (question, accepted answers, documents that support the answer)
QUESTIONS = [
    ("Who leads the team that owns the Atlas scheduler?", ["Priya Raman"], ["d1", "d2"]),
    ("Who leads the team that owns Beacon?", ["Tomas Ekberg"], ["d3", "d4"]),
    ("How many times does Atlas retry a failed job?", ["three", "3"], ["d8"]),
    ("Which company builds the charging docks in the office where the Fleet Platform team sits?", ["Voltline"], ["d2", "d5"]),
    ("How many days of data does the telemetry pipeline keep?", ["30"], ["d3"]),
]

"""Keyword retrieval for the agent's search tool."""


def tokenize(text):
    """Lowercase words of letters and digits."""
    raise NotImplementedError


class BM25:
    def __init__(self, docs, k1=1.5, b=0.75):
        """docs: {doc_id: text}."""
        raise NotImplementedError

    def search(self, query, k=3):
        """The k best (doc_id, score) pairs with score > 0, best first;
        ties broken by doc_id."""
        raise NotImplementedError


def recall_at_k(index, questions, k=3):
    """Share of supporting documents found by searching each question text
    directly. questions: [(question, answers, supporting_doc_ids)]."""
    raise NotImplementedError

"""The agent loop. The model is any callable that takes the message list and
returns one action:

    {"search": "<query>"}                                  # call the search tool
    {"answer": "<text>", "citations": ["d1", "d2"]}        # finish

The real take-home wraps a hosted model with tool calling behind this same
interface; tests pass a scripted model.
"""


def run_agent(question, model, search, max_steps=6, k=3):
    """Run until the model answers or max_steps model calls have been made.

    search(query, k) returns [(doc_id, text)]. Every search result goes back to
    the model as a tool message. Returns a dict:
      answer       the model's answer, or None if the budget ran out
      citations    cited doc ids that the agent actually retrieved
      unsupported  cited doc ids it never retrieved (a hallucination signal)
      steps        model calls made
      stopped      "answer" or "budget"
    A repeated query (same words, ignoring case and spacing) is not run again;
    the model is told so instead.
    """
    raise NotImplementedError


def evaluate(results, questions):
    """results[i] is run_agent's output for questions[i] = (question, answers,
    supporting_doc_ids). Returns {"accuracy", "citation_recall",
    "unsupported_rate", "mean_steps"}: accuracy counts an answer that contains
    any accepted answer (case-insensitive); citation recall is the share of
    supporting docs cited; unsupported rate is the share of results with any
    unsupported citation."""
    raise NotImplementedError