Long Context vs RAG: Accuracy, Cost and Latency on the Same 24 Questions

"""Long context or RAG for Harbourline Freight's depot handbook. You write the functions marked TODO, one step at
a time."""
import re
from concurrent.futures import ThreadPoolExecutor

import numpy as np

import harness as H

RULES = """Answer from the handbook text only, in as few words as possible.
Amendments in Part 3 replace anything earlier that they contradict.
For a list, give every item that qualifies, separated by commas, and nothing else.
If the handbook does not give the answer, reply exactly: not stated"""


# ---------- Step 1: the whole handbook in the prompt ----------

def long_prompt(handbook, question):
    """The whole handbook, then RULES, then the question."""
    # TODO (Step 1): the handbook, then RULES, then the question.
    raise NotImplementedError("Step 1: write long_prompt()")


def mentions(text, phrase):
    """True if phrase appears in text as a whole word or number, ignoring case."""
    # TODO (Step 1): re.search with (?<![\w.]) before re.escape(phrase.lower()) and (?![\w]) after, on text.lower().
    raise NotImplementedError("Step 1: write mentions()")


def grade(answer, q):
    """True if the answer mentions everything in q["must"] and nothing in q["must_not"]."""
    # TODO (Step 1): every q["must"] mentioned, and no q["must_not"] mentioned.
    raise NotImplementedError("Step 1: write grade()")


def run_all(answer_for, questions, workers=4):
    """For each question, the reply answer_for(question text) plus "type" and "correct" (grade()), in question order,
    `workers` questions at a time."""
    # TODO (Step 1): ThreadPoolExecutor(workers).map(answer_for, question texts), then add type and correct.
    raise NotImplementedError("Step 1: write run_all()")


# ---------- Step 2: what each answer costs ----------

def cost_usd(result):
    """Dollars for one call at H.PRICE (per million tokens, input and output priced separately)."""
    raise NotImplementedError("cost_usd() arrives in Step 2")


def percentile(values, p):
    """Nearest-rank percentile: the smallest value with at least p% of the values at or below it."""
    raise NotImplementedError("percentile() arrives in Step 2")


def summary(results):
    """{"accuracy", "prompt_tokens" (mean), "usd_per_1000" (questions), "p50_s", "p95_s", "by_type": {type: accuracy}}."""
    raise NotImplementedError("summary() arrives in Step 2")


# ---------- Step 3: retrieve sections instead ----------

def split_sections(handbook):
    """[{"title", "text"}] for every "## " section: the title without "## ", the text from its heading line up to
    the next heading of any level ("# " or "## "), stripped."""
    raise NotImplementedError("split_sections() arrives in Step 3")


def retrieve(question, sections, k):
    """The k sections whose text embeddings are closest to the question's, closest first."""
    raise NotImplementedError("retrieve() arrives in Step 3")


def rag_prompt(question, chunks):
    """The chunks' texts (separated by blank lines), then RULES, then the question."""
    raise NotImplementedError("rag_prompt() arrives in Step 3")


# ---------- Step 4: amendments travel with their section ----------

def with_amendments(chunks, sections):
    """The chunks, each followed by every amendment whose title ends with "to section <that chunk's number>" (the
    number is the first word of a title, such as 2.14), with no section twice."""
    raise NotImplementedError("with_amendments() arrives in Step 4")


# ---------- Step 5: list questions, a few profiles at a time ----------

def is_list_question(question):
    """True for questions that ask which depots (plural) qualify: they start with "Which depots", any case."""
    raise NotImplementedError("is_list_question() arrives in Step 5")


def map_prompt(question, group, sections):
    """The profiles in group, each with its amendments, then: for each depot, one line "<name>: <deciding fact> =>
    yes" or "... => no", saying whether that depot alone qualifies for the question."""
    raise NotImplementedError("map_prompt() arrives in Step 5")


def map_answer(question, sections, batch=6, workers=4):
    """Ask map_prompt() about the depot profiles `batch` at a time and join the depots whose line ends "=> yes". Returns a
    reply like H.ask(): "text" is the names joined with ", " ("none" if there are none), tokens are summed over the
    calls, and "seconds" is the latency with `workers` calls at a time: the sum, over each round of `workers` calls,
    of the slowest call in it."""
    raise NotImplementedError("map_answer() arrives in Step 5")


def routed(question, sections, k=4):
    """map_answer() for list questions; otherwise RAG with retrieve(k) and with_amendments()."""
    raise NotImplementedError("routed() arrives in Step 5")