RAG Pipeline Walkthrough

RAG Pipeline Walkthrough#

Notebook 01 described the data; notebook 03 reports the benchmark numbers. This middle notebook explains how the pipeline itself is wired together – function by function, with small concrete inputs and outputs – so that when the results show up in notebook 03 you already know exactly what was done to produce each number.

Everything here is imported from src/rag_benchmark.py and src/embeddings_index.py. The runnable scripts call OpenAI; this notebook stays mostly offline and uses saved CSV outputs from a previous run to show what the model actually returned. That way the walkthrough is cheap to re-execute and does not require an API key (it does briefly call the embeddings API in Step 3b to embed one query, but that costs a fraction of a cent and can be skipped).

Two flavors of retrieval#

“RAG” just means retrieve text from a corpus, then generate. It does not specify how retrieval works. We benchmark two retrieval families in this homework:

  • Structured retrieval (find_10k_filing): the question already carries a CIK and a fiscal-year-end, so we look up the matching 10-K directly and run a keyword-based chunk extractor over it. This is what the rag_cleaned and rag_raw conditions use.

  • Dense semantic retrieval (semantic_search): the corpus is pre-chunked and every chunk is embedded into a 1536-dim vector with the OpenAI embeddings API. At query time we embed the question and ask FAISS for the top-k nearest chunks across the whole corpus – no CIK/year prefilter. This is what the rag_semantic_cleaned condition uses, and it is what most people mean when they say “classical RAG.”

Both end up calling the same downstream code (build_rag_prompt, run_prompt, extract_number, is_correct). Only the retrieval step changes.

The pipeline at a glance#

For each (company, fiscal year, metric) question we do the same thing in four conditions (baseline, RAG-cleaned, RAG-raw, RAG-semantic-cleaned):

    build_test_cases()              # generate questions from Compustat
            |
            v
    ---- pick a retrieval path (skipped in baseline) ----
      structured:  find_10k_filing()  +  extract_financial_context_*()
      semantic:    semantic_search()  over a prebuilt FAISS index
            |
            v
    build_rag_prompt()              # assemble the prompt
            |
            v
    run_prompt()                    # OpenAI chat.completions call
            |
            v
    extract_number()                # parse a number out of the reply
            |
            v
    is_correct()                    # compare to Compustat ground truth

The sections below walk through those steps one at a time.

from pathlib import Path

import pandas as pd

from settings import config
from rag_benchmark import (
    MODEL_ID,
    MODEL_INPUT_PRICE,
    MODEL_OUTPUT_PRICE,
    SYSTEM_MSG,
    METRICS,
    MAX_WORDS_CLEANED,
    MAX_WORDS_RAW,
    FINANCIAL_KEYWORDS,
    build_test_cases,
    extract_financial_context_cleaned,
    extract_financial_context_raw,
    load_filing_index,
    find_10k_filing,
    build_rag_prompt,
    extract_number,
    is_correct,
    calculate_cost,
)
from embeddings_index import (
    EMBED_MODEL,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    chunk_text,
    load_index,
    semantic_search,
)

DATA_DIR = Path(config("DATA_DIR"))
OUTPUT_DIR = Path(config("OUTPUT_DIR"))

Step 1: generate benchmark questions#

build_test_cases(min_year=...) reads Compustat, filters to recent fiscal years, and produces one dictionary per (company, fiscal_year, metric) combination. The three metrics are fixed in rag_benchmark.py:

METRICS

{'sale': 'total revenue (net sales)',
 'ni': 'net income (loss)',
 'at': 'total assets'}

Each returned test case is a plain dict. The LLM will see question, the grader will compare its answer against expected, and the retriever keys off cik and fiscal_year_end.

test_cases = build_test_cases(min_year=2023)
print(f"Total questions: {len(test_cases)}")
print(f"  = {len({tc['cik'] for tc in test_cases})} companies "
      f"x {len({tc['fiscal_year'] for tc in test_cases})} fiscal years "
      f"x {len(METRICS)} metrics "
      "(minus any missing Compustat rows)")

sample_case = test_cases[0]
sample_case

Total questions: 96
  = 10 companies x 4 fiscal years x 3 metrics (minus any missing Compustat rows)

{'company_name': 'APPLE INC',
 'cik': '0000320193',
 'fiscal_year': 2023,
 'fiscal_year_end': '2023-09-30',
 'metric': 'sale',
 'metric_label': 'total revenue (net sales)',
 'expected': 383285.0,
 'question': "What was APPLE INC's total revenue (net sales) for the fiscal year ending 2023-09-30? Report a single number in millions of US dollars."}

The natural-language question assembled for the LLM:

print(sample_case["question"])

What was APPLE INC's total revenue (net sales) for the fiscal year ending 2023-09-30? Report a single number in millions of US dollars.

Step 2: find the matching 10-K#

Retrieval here is a structured lookup, not semantic search. We have a filing index keyed by CIK, and for each question we keep the first 10-K whose filing date falls in the 180 days after the fiscal year end. That window is wide enough to catch the annual report but narrow enough to avoid grabbing the next year’s filing.

tenk_idx = load_filing_index()
print(f"10-K filings in the index: {len(tenk_idx)}")
tenk_idx[["cik", "coname", "form", "fdate"]].sort_values("fdate").head()

10-K filings in the index: 48
cik coname form fdate
47 0001065280 NETFLIX INC 10-K 2022-01-27
46 0001652044 Alphabet Inc. 10-K 2022-02-02
45 0001018724 AMAZON COM INC 10-K 2022-02-04
44 0001318605 Tesla, Inc. 10-K 2022-02-07
43 0000200406 JOHNSON & JOHNSON 10-K 2022-02-17

find_10k_filing takes the index, a CIK, a fiscal-year-end date, a directory of filings on disk, and a chunk-extractor function. It returns the extracted context plus a bit of provenance, or None if the filing is missing from our corpus.

fc_clean = find_10k_filing(
    tenk_idx,
    cik=sample_case["cik"],
    fiscal_year_end=sample_case["fiscal_year_end"],
    filings_dir=DATA_DIR / "wrds_clean_filings",
    extractor=extract_financial_context_cleaned,
)
{k: (len(v) if k == "context" else v) for k, v in fc_clean.items()}

{'context': 18868,
 'company_name': 'Apple Inc.',
 'filing_date': '2023-11-03',
 'filing_bytes': 1190332}

The same call against the raw archive – same CIK, same date, but pointing at wrds_raw_filings and using the raw extractor. Notice the context length jumps because the raw text carries markup.

fc_raw = find_10k_filing(
    tenk_idx,
    cik=sample_case["cik"],
    fiscal_year_end=sample_case["fiscal_year_end"],
    filings_dir=DATA_DIR / "wrds_raw_filings",
    extractor=extract_financial_context_raw,
)
{k: (len(v) if k == "context" else v) for k, v in fc_raw.items()}

{'context': 72546,
 'company_name': 'Apple Inc.',
 'filing_date': '2023-11-03',
 'filing_bytes': 9569569}

Step 3: chunk selection#

extract_financial_context_cleaned is the real workhorse of the retrieval step. It scans the 10-K for two kinds of content:

  1. Keyword + number lines (e.g. “Total revenue $383,285”) and grabs a few lines of surrounding context.

  2. Financial narrative paragraphs that mention at least two financial keywords and contain a dollar figure.

The keyword list is the same set the LLM is trying to answer about:

FINANCIAL_KEYWORDS

['net sales',
 'total net revenue',
 'total revenue',
 'net income',
 'net loss',
 'net earnings',
 'total assets',
 'results of operations',
 'operating income',
 'gross margin',
 'gross profit']

Both conditions (cleaned and raw) use the same extractor logic – extract_financial_context_raw is deliberately a thin wrapper. The only difference between conditions is the text WRDS hands us. A toy example makes the behavior concrete:

toy_10k = """
ITEM 7. RESULTS OF OPERATIONS

The following discussion compares our fiscal 2024 results to fiscal 2023.

Net sales increased 8% to $383,285 million, driven by higher iPhone and
Services revenue.

Total assets grew to $352,755 million as of September 28, 2024.

Unrelated filler about forward-looking statements goes here and should
not be picked up by the chunker because it has neither a keyword nor a
dollar figure.

Net income for fiscal 2024 was $96,995 million, compared to $99,803
million in the prior year.
"""
print(extract_financial_context_cleaned(toy_10k, max_words=500))

ITEM 7. RESULTS OF OPERATIONS

The following discussion compares our fiscal 2024 results to fiscal 2023.

Net sales increased 8% to $383,285 million, driven by higher iPhone and
Services revenue.

Total assets grew to $352,755 million as of September 28, 2024.

Unrelated filler about forward-looking statements goes here and should

---

The following discussion compares our fiscal 2024 results to fiscal 2023.

Net sales increased 8% to $383,285 million, driven by higher iPhone and
Services revenue.

Total assets grew to $352,755 million as of September 28, 2024.

Unrelated filler about forward-looking statements goes here and should
not be picked up by the chunker because it has neither a keyword nor a
dollar figure.

---

Unrelated filler about forward-looking statements goes here and should
not be picked up by the chunker because it has neither a keyword nor a
dollar figure.

Net income for fiscal 2024 was $96,995 million, compared to $99,803
million in the prior year.

Even on real filings, the extracted context is a small fraction of the full 10-K. Here is the cleaned vs raw comparison for our sample filing:

cleaned_words = len(fc_clean["context"].split())
raw_words = len(fc_raw["context"].split())
pd.DataFrame({
    "condition":     ["rag_cleaned", "rag_raw"],
    "budget_words":  [MAX_WORDS_CLEANED, MAX_WORDS_RAW],
    "actual_words":  [cleaned_words, raw_words],
    "chars":         [len(fc_clean["context"]), len(fc_raw["context"])],
    "chars_per_word": [
        round(len(fc_clean["context"]) / max(cleaned_words, 1), 1),
        round(len(fc_raw["context"]) / max(raw_words, 1), 1),
    ],
})
condition budget_words actual_words chars chars_per_word
0 rag_cleaned 3000 3009 18868 6.3
1 rag_raw 3000 2996 72546 24.2

The raw chunks have far more characters per word – that’s the SGML/HTML/XBRL markup that WRDS’s cleaning pipeline strips. Those characters become input tokens, and input tokens cost money.

First 500 characters of the two contexts so you can see the difference directly:

print("--- cleaned (first 500 chars) ---")
print(fc_clean["context"][:500])
print("\n--- raw (first 500 chars) ---")
print(fc_raw["context"][:500])

--- cleaned (first 500 chars) ---
Form 10-K Summary 

57

This Annual Report on Form 10-K (Form 10-K) contains forward-looking statements, within the meaning of the Private Securities Litigation Reform Act of 1995, that involve risks and uncertainties. Many of the forward-looking statements are located in Part I, Item 1 of this Form 10-K under the heading Business and Part II, Item 7 of this Form 10-K under the heading Managements Discussion and Analysis of Financial Condition and Results of Operations. Forward-looking statement

--- raw (first 500 chars) ---
<DESCRIPTION>EX-31.1
<TEXT>
<html><head>
<!-- Document created using Wdesk -->
<!-- Copyright 2023 Workiva -->
<title>Document</title></head><body><div id="i2369e0a153494a43a2376c67a756c66b_1"></div><div style="min-height:42.75pt;width:100%"><div><font><br></font></div></div><div style="text-align:right"><font style="color:#000000;font-family:'Helvetica',sans-serif;font-size:9pt;font-weight:700;line-height:120%">Exhibit 31.1</font></div><div style="margin-top:12pt;text-align:center"><font style=

Step 3b: dense semantic retrieval (the “classical RAG” path)#

The structured-lookup retrieval above works because the question already names the company and the fiscal year, and we have a clean index keyed on those. Most production RAG stacks do not have that luxury – the question is free text and there is no obvious join key. So they fall back to semantic search: embed the corpus once, embed the query, and ask a vector index for the most similar chunks.

Three pieces are involved:

  1. A chunker that splits each filing into overlapping passages. chunk_text(text, chunk_size, overlap) is a sliding character window – nothing fancy, just chunk_size characters at a time with overlap characters of overlap so a sentence isn’t sliced in half.

  2. An embedding model that turns each chunk (and later, each query) into a fixed-length vector. We use OpenAI text-embedding-3-small – 1536 dimensions, same API client we already use for the chat completions.

  3. A vector index that, given a query vector, returns the top-k most similar chunk vectors. We use FAISS (IndexFlatIP with L2-normalized vectors, which gives exact cosine similarity over the ~55k chunks in the cleaned corpus).

The build (task_build_embeddings in dodo.py -> python src/embeddings_index.py) chunks all 48 cleaned filings, embeds every chunk in batches, and writes:

  • _data/embeddings/cleaned_chunks.parquet – one row per chunk with metadata (cik, accession, fdate, char_start, char_end, text).

  • _data/embeddings/cleaned_index.faiss – the FAISS index, with row order aligned to the parquet.

We do this once and cache, so subsequent benchmark runs only pay the query-side embedding cost.

Chunking demo#

Same toy 10-K we used above for the keyword extractor, this time fed to chunk_text:

toy_chunks = chunk_text(toy_10k, chunk_size=200, overlap=40)
pd.DataFrame([
    {
        "chunk_idx": i,
        "char_start": s,
        "char_end": e,
        "preview": (t[:80] + "...") if len(t) > 80 else t,
    }
    for i, (t, s, e) in enumerate(toy_chunks)
])
chunk_idx char_start char_end preview
0 0 0 200 \nITEM 7. RESULTS OF OPERATIONS\n\nThe followi...
1 1 160 360 higher iPhone and\nServices revenue.\n\nTotal...
2 2 320 516 re and should\nnot be picked up by the chunker...

In the real build we use chunk_size=2000, overlap=200. Smaller chunks here just so the demo fits on a screen. The values used by the build are imported from embeddings_index.py:

print(f"CHUNK_SIZE   = {CHUNK_SIZE}")
print(f"CHUNK_OVERLAP = {CHUNK_OVERLAP}")
print(f"EMBED_MODEL  = {EMBED_MODEL}")

CHUNK_SIZE   = 2000
CHUNK_OVERLAP = 200
EMBED_MODEL  = text-embedding-3-small

Loading the prebuilt index#

Once the build is done, querying is fast and offline (modulo a single embedding call per query). Each chunk row carries its source filing’s CIK and filing date, so we can tell after retrieval which company / year the chunk came from.

chunks_df, faiss_index = load_index(
    DATA_DIR / "embeddings" / "cleaned_chunks.parquet",
    DATA_DIR / "embeddings" / "cleaned_index.faiss",
)
print(f"chunks in index: {len(chunks_df):,}")
print(f"distinct filings covered: {chunks_df['accession'].nunique()}")
print(f"vector dimension: {faiss_index.d}")
chunks_df.head(3)[["coname", "fdate", "chunk_idx", "char_start", "char_end"]]

chunks in index: 63,252
distinct filings covered: 48
vector dimension: 1536
coname fdate chunk_idx char_start char_end
0 Walmart Inc. 2026-03-13 0 0 2000
1 Walmart Inc. 2026-03-13 1 1800 3800
2 Walmart Inc. 2026-03-13 2 3600 5600

A semantic-search query#

semantic_search(client, question, chunks_df, faiss_index, top_k, max_words) embeds the question, calls FAISS for the top-k chunks, and concatenates them up to a word budget. It returns the same shape of dict as find_10k_filing plus a top_k_filings field that lets us inspect which filings the chunks came from.

We need an OpenAI client only to embed the question (the corpus embeddings are already cached). If you do not have an API key handy, skip the next two cells and read the prose – the rest of the notebook does not depend on this call.

api_key = config("OPENAI_API_KEY", default=None)
if api_key:
    from openai import OpenAI
    client = OpenAI(api_key=api_key)
    sem_result = semantic_search(
        client,
        sample_case["question"],
        chunks_df,
        faiss_index,
        top_k=5,
        max_words=600,
    )
else:
    client = None
    sem_result = None
    print("Set OPENAI_API_KEY to run the live retrieval cell below.")

sem_result is None or pd.DataFrame(sem_result["top_k_filings"])
rank cik coname fdate score
0 0 0000320193 Apple Inc. 2022-10-28 0.696436
1 1 0000320193 Apple Inc. 2022-10-28 0.690374
2 2 0000320193 Apple Inc. 2024-11-01 0.686319
3 3 0000320193 Apple Inc. 2023-11-03 0.679418
4 4 0000320193 Apple Inc. 2023-11-03 0.679405

Read across the rows: each row is one of the top-k chunks the semantic retriever picked, and the coname / fdate columns tell us whether the retriever is finding chunks from the right filing (the question asks about sample_case["company_name"] for fiscal year sample_case["fiscal_year_end"]).

Keyword vs semantic: side by side#

Same question, two retrievers. The keyword extractor pulls passages only from the matched filing; the semantic retriever pulls from anywhere in the corpus, but tends to land on the right company when the question is specific enough.

if sem_result is not None:
    sem_filings = sorted({r["coname"] for r in sem_result["top_k_filings"]})
    sem_chars = len(sem_result["context"])
    sem_words = len(sem_result["context"].split())
else:
    sem_filings, sem_chars, sem_words = ["(skipped)"], 0, 0

pd.DataFrame({
    "retriever":      ["keyword (structured)", "semantic (global FAISS)"],
    "filings_used":   [
        fc_clean["company_name"],
        ", ".join(sem_filings),
    ],
    "context_chars":  [len(fc_clean["context"]), sem_chars],
    "context_words":  [len(fc_clean["context"].split()), sem_words],
})
retriever filings_used context_chars context_words
0 keyword (structured) Apple Inc. 18868 3009
1 semantic (global FAISS) Apple Inc. 2060 340

Step 4: assemble the prompt#

build_rag_prompt wraps the retrieved context and the question into a single user message. When filing_context is None (baseline or missing filing) it falls back to sending the bare question – this is how the baseline condition and “filing not found” cases are handled with the same downstream code path.

print("--- BASELINE PROMPT (no context) ---")
print(build_rag_prompt(sample_case["question"], None))

--- BASELINE PROMPT (no context) ---
What was APPLE INC's total revenue (net sales) for the fiscal year ending 2023-09-30? Report a single number in millions of US dollars.

print("--- RAG PROMPT (first 1200 chars) ---")
rag_prompt = build_rag_prompt(sample_case["question"], fc_clean)
print(rag_prompt[:1200])
print("...")
print(f"\n[total length: {len(rag_prompt):,} chars]")

--- RAG PROMPT (first 1200 chars) ---
Use the following SEC filing excerpts to answer the question. Extract the specific numbers from the excerpts.

SEC FILING EXCERPTS (Source: Apple Inc., 10-K, filed 2023-11-03):
Form 10-K Summary 

57

This Annual Report on Form 10-K (Form 10-K) contains forward-looking statements, within the meaning of the Private Securities Litigation Reform Act of 1995, that involve risks and uncertainties. Many of the forward-looking statements are located in Part I, Item 1 of this Form 10-K under the heading Business and Part II, Item 7 of this Form 10-K under the heading Managements Discussion and Analysis of Financial Condition and Results of Operations. Forward-looking statements provide current expectations of future events based on certain assumptions and include any statement that does not directly relate to any historical or current fact. For example, statements in this Form 10-K regarding the potential future impact of macroeconomic conditions on the Companys business and results of operations are forward-looking statements. Forward-looking statements can also be identified by words such as future, anticipates, believes, estimates, expects, intends, plans, predicts, will, would, could, 
...

[total length: 19,248 chars]

The system message that accompanies every call is also fixed:

print(SYSTEM_MSG)

You are a financial analyst. Answer the question with a single numeric value in millions of US dollars. If you are uncertain, give your best estimate. Do not refuse to answer.

Step 5: the LLM call#

run_prompt(client, model_id, prompt, system_msg=...) is a thin wrapper around openai.chat.completions.create. It measures wall time, captures token counts, and returns a dict:

{
    "output":        "...",   # the model's reply text
    "latency_s":     0.87,
    "input_tokens":  1234,
    "output_tokens": 12,
}

We do not call the API here – that happens in src/rag_benchmark.py at build time and the results are cached to CSV. Let’s load the cached results and look at one real reply for the sample question.

df_baseline = pd.read_csv(OUTPUT_DIR / "baseline_results.csv")
df_clean    = pd.read_csv(OUTPUT_DIR / "rag_cleaned_results.csv")
df_raw      = pd.read_csv(OUTPUT_DIR / "rag_raw_results.csv")
df_semantic = pd.read_csv(OUTPUT_DIR / "rag_semantic_cleaned_results.csv")

sample_cik_int = int(sample_case["cik"])

def pick(df):
    m = (
        (df["cik"].astype(int) == sample_cik_int)
        & (df["fiscal_year"] == sample_case["fiscal_year"])
        & (df["metric"] == sample_case["metric"])
    )
    return df.loc[m].iloc[0]

row_b = pick(df_baseline)
row_c = pick(df_clean)
row_r = pick(df_raw)
row_s = pick(df_semantic)

pd.DataFrame({
    "condition":    ["baseline", "rag_cleaned", "rag_raw", "rag_semantic_cleaned"],
    "answer":       [row_b["answer"], row_c["answer"], row_r["answer"], row_s["answer"]],
    "latency_s":    [row_b["latency_s"], row_c["latency_s"], row_r["latency_s"], row_s["latency_s"]],
    "input_tokens": [row_b["input_tokens"], row_c["input_tokens"], row_r["input_tokens"], row_s["input_tokens"]],
    "output_tokens":[row_b["output_tokens"], row_c["output_tokens"], row_r["output_tokens"], row_s["output_tokens"]],
})
condition answer latency_s input_tokens output_tokens
0 baseline Apple Inc's total revenue for the fiscal year ... 2.827 82 30
1 rag_cleaned 383,285 1.223 4227 3
2 rag_raw 394328 million 1.886 21908 3
3 rag_semantic_cleaned 383,285 4.520 5404 3

Input-token counts line up with intuition: baseline is tiny (just the question), RAG-cleaned is larger (question + chunked context), and RAG-raw is larger still (same chunking, but each chunk carries markup).

Step 6: parse a number out of the reply#

The LLM answers in prose (“Apple’s fiscal 2024 revenue was approximately $383 billion.”) but the grader needs a numeric value in millions of USD. extract_number handles the common shapes:

examples = [
    "Apple's fiscal 2024 revenue was approximately $383 billion.",
    "Net income was 96,995 million.",
    "The total assets figure is $352,755.",
    "Approximately -2,000 million (net loss).",
    "I don't know.",
]
pd.DataFrame({
    "llm_reply": examples,
    "extracted_millions": [extract_number(x) for x in examples],
})
llm_reply extracted_millions
0 Apple's fiscal 2024 revenue was approximately ... 383000.0
1 Net income was 96,995 million. 96995.0
2 The total assets figure is $352,755. 352755.0
3 Approximately -2,000 million (net loss). 2000.0
4 I don't know. NaN

The parser’s preference order is:

  1. "X billion" -> X * 1000 (convert to millions).

  2. "X million" -> X.

  3. Otherwise, fall back to the last numeric token in the string (typically the figure the model quoted last).

It’s deliberately forgiving – these are small models and the prose varies. The 5% and 10% tolerance bands in the next step absorb any rounding.

Step 7: grade the answer#

is_correct(expected, extracted, tolerance) is a relative error check with a guard for zero expected values.

expected = sample_case["expected"]
examples = [expected, expected * 1.03, expected * 1.08, expected * 0.5, None]
pd.DataFrame({
    "expected": [expected] * len(examples),
    "extracted": examples,
    "within_5pct": [is_correct(expected, x, 0.05) for x in examples],
    "within_10pct": [is_correct(expected, x, 0.10) for x in examples],
})
expected extracted within_5pct within_10pct
0 383285.0 383285.00 True True
1 383285.0 394783.55 True True
2 383285.0 413947.80 False True
3 383285.0 191642.50 False False
4 383285.0 NaN False False

Notebook 03 reports accuracy at both tolerance levels because a fair number of correct answers in this task miss 5% by rounding and land comfortably inside 10%.

Step 8: cost accounting#

calculate_cost(input_tokens, output_tokens) applies the published gpt-4o-mini prices (per million tokens):

print(f"Model:              {MODEL_ID}")
print(f"Input price (USD/M tokens):  {MODEL_INPUT_PRICE}")
print(f"Output price (USD/M tokens): {MODEL_OUTPUT_PRICE}")

pd.DataFrame({
    "condition":     ["baseline", "rag_cleaned", "rag_raw", "rag_semantic_cleaned"],
    "input_tokens":  [row_b["input_tokens"], row_c["input_tokens"], row_r["input_tokens"], row_s["input_tokens"]],
    "output_tokens": [row_b["output_tokens"], row_c["output_tokens"], row_r["output_tokens"], row_s["output_tokens"]],
    "usd_per_query": [
        calculate_cost(row_b["input_tokens"], row_b["output_tokens"]),
        calculate_cost(row_c["input_tokens"], row_c["output_tokens"]),
        calculate_cost(row_r["input_tokens"], row_r["output_tokens"]),
        calculate_cost(row_s["input_tokens"], row_s["output_tokens"]),
    ],
})

Model:              gpt-4o-mini
Input price (USD/M tokens):  0.15
Output price (USD/M tokens): 0.6
condition input_tokens output_tokens usd_per_query
0 baseline 82 30 0.000030
1 rag_cleaned 4227 3 0.000636
2 rag_raw 21908 3 0.003288
3 rag_semantic_cleaned 5404 3 0.000812

One RAG call on this sample question costs well under a cent – which is the whole point of the exercise: if a small, cheap model can read the 10-K, you don’t need an expensive one to know the answer.

Putting it together end-to-end#

Here is what the pipeline does for our one sample question, traced through every step.

trace = {
    "1. test case":      f"{sample_case['company_name']} / "
                         f"{sample_case['metric_label']} / "
                         f"FY ending {sample_case['fiscal_year_end']}",
    "2. ground truth":   f"{sample_case['expected']:,.0f} (USD millions, from Compustat)",
    "3. filing found":   f"{fc_clean['company_name']}, "
                         f"filed {fc_clean['filing_date']}, "
                         f"{fc_clean['filing_bytes']:,} bytes on disk",
    "4. context chars":  f"cleaned={len(fc_clean['context']):,}  "
                         f"raw={len(fc_raw['context']):,}  "
                         f"semantic={row_s['context_chars']:,}",
    "5. llm reply (rag_cleaned)":   row_c["answer"][:200],
    "5. llm reply (rag_semantic)":  row_s["answer"][:200],
    "6. extracted (cleaned)":       f"{row_c['extracted']}",
    "6. extracted (semantic)":      f"{row_s['extracted']}",
    "7. correct @ 10% (cleaned)":   bool(row_c["correct_10pct"]),
    "7. correct @ 10% (semantic)":  bool(row_s["correct_10pct"]),
    "7. semantic top-1 chunk match": bool(row_s["top_chunk_cik_match"]),
    "8. cost cleaned (USD)":   calculate_cost(row_c["input_tokens"], row_c["output_tokens"]),
    "8. cost semantic (USD)":  calculate_cost(row_s["input_tokens"], row_s["output_tokens"]),
}
for k, v in trace.items():
    print(f"{k:>32}: {v}")

                    1. test case: APPLE INC / total revenue (net sales) / FY ending 2023-09-30
                 2. ground truth: 383,285 (USD millions, from Compustat)
                 3. filing found: Apple Inc., filed 2023-11-03, 1,190,332 bytes on disk
                4. context chars: cleaned=18,868  raw=72,546  semantic=18,596
      5. llm reply (rag_cleaned): 383,285
     5. llm reply (rag_semantic): 383,285
          6. extracted (cleaned): 383285.0
         6. extracted (semantic): 383285.0
      7. correct @ 10% (cleaned): True
     7. correct @ 10% (semantic): True
   7. semantic top-1 chunk match: True
           8. cost cleaned (USD): 0.000636
          8. cost semantic (USD): 0.000812

Where to go from here#

Notebook 03 scales this up: it runs the full grid of questions through all four conditions and reports aggregate accuracy, latency, token usage, and cost. Before opening it, you should now be able to predict the qualitative shape of the results:

  • Baseline is cheap but inaccurate (the model is guessing).

  • RAG-cleaned (structured) is accurate and still cheap (small model, compact context, clean text, exact filing match).

  • RAG-raw (structured) is accurate-ish but costs more per query and sometimes trips on markup – WRDS’s cleaning pipeline earns its keep.

  • RAG-semantic-cleaned trades the exact filing match for a global FAISS lookup. Whether this helps or hurts depends on whether the top-k chunks land on the right company; top_chunk_cik_match in the per-question CSV tells you.

Levers you can pull:

  • Change FINANCIAL_KEYWORDS / MAX_WORDS_* in src/rag_benchmark.py to reshape the keyword chunker.

  • Change CHUNK_SIZE / CHUNK_OVERLAP / EMBED_MODEL in src/embeddings_index.py (then rerun task_build_embeddings) to reshape the semantic chunker.

  • Change top_k / max_words in run_rag_semantic to give the LLM more or fewer retrieved chunks per question.

Contents