primer.common.corpus

  1"""
  2A small enterprise knowledge base used by the embeddings and RAG examples.
  3
  4It's deliberately built to exercise the failure modes that bite real retrieval systems:
  5
  6* **Hard negatives.** `it-001` (how to *reset* your password) and `it-002`
  7  (password *policy*) share a topic but answer different questions.
  8  Topic-matching retrievers confuse them.
  9* **Exact identifiers.** `it-004` is about error `ERR-4012`, and `it-005`
 10  about `ERR-4013`. Dense embeddings blur these; BM25 nails them.
 11* **Synonyms.** A query about "automobile reimbursement" should find the
 12  "car mileage expense" doc with no shared words. BM25 misses it; dense
 13  retrieval finds it.
 14* **Access control.** Some docs carry `acl` groups (e.g. only `finance` or
 15  `hr` may see them) for permission-aware retrieval in `primer.agents.rag`.
 16* **Freshness.** `updated` dates let you demo metadata filtering and stale
 17  content (e.g. two versions of the travel policy).
 18
 19`LABELED_QUERIES` maps realistic questions to the doc IDs that answer them,
 20the "golden set" for recall@k / MRR / nDCG in
 21`primer.ml.metrics` and `primer.ml.embeddings.retrieval`.
 22"""
 23
 24from __future__ import annotations
 25
 26from dataclasses import dataclass, field
 27
 28
 29@dataclass(frozen=True)
 30class Doc:
 31    id: str
 32    title: str
 33    text: str
 34    department: str
 35    updated: str  # ISO date
 36    acl: frozenset[str] = field(default_factory=lambda: frozenset({"everyone"}))
 37
 38
 39DOCS: list[Doc] = [
 40    Doc("it-001", "How to reset your password",
 41        "If you forgot your password or your account is locked, go to the self-service portal, "
 42        "choose 'Forgot password', verify with your authenticator app, and set a new password. "
 43        "The reset link expires after 15 minutes.",
 44        "IT", "2026-03-02"),
 45    Doc("it-002", "Password policy",
 46        "Passwords must be at least 14 characters and include a number and a symbol. "
 47        "Passwords expire every 180 days and the last 10 passwords cannot be reused. "
 48        "This policy applies to all employees and contractors.",
 49        "IT", "2025-11-20"),
 50    Doc("it-003", "Setting up the VPN",
 51        "Install the AnyConnect client from the software center, sign in with your company credentials, "
 52        "and approve the two-factor prompt. Use the VPN for all remote access to internal systems.",
 53        "IT", "2026-01-15"),
 54    Doc("it-004", "Error ERR-4012: VPN tunnel failed",
 55        "ERR-4012 means the VPN tunnel could not be established, usually because the client is out of date. "
 56        "Update AnyConnect to version 5.1 or later and reboot.",
 57        "IT", "2026-04-10"),
 58    Doc("it-005", "Error ERR-4013: certificate expired",
 59        "ERR-4013 means your device certificate expired. Open the software center and run "
 60        "'Renew device certificate', then reconnect.",
 61        "IT", "2026-04-10"),
 62    Doc("it-006", "Setting up two-factor authentication",
 63        "Download the authenticator app, scan the QR code on the security page, and enter the six digit code "
 64        "to finish enrolling in MFA. Two-factor is required for email and VPN.",
 65        "IT", "2025-09-01"),
 66    Doc("it-007", "Reporting phishing emails",
 67        "If an email looks suspicious, do not click links. Use the 'Report phishing' button in Outlook. "
 68        "Security reviews every report within one business day.",
 69        "IT", "2026-02-11"),
 70    Doc("it-008", "Requesting a new laptop",
 71        "Laptops are refreshed every three years. Submit a hardware request in the IT portal with your "
 72        "manager's approval. Standard devices are MacBook Pro or ThinkPad X1.",
 73        "IT", "2025-12-05"),
 74    Doc("it-009", "Printer troubleshooting",
 75        "If printing fails, check the toner and paper tray, then remove and re-add the printer from settings. "
 76        "Floor printers are named by building and floor.",
 77        "IT", "2024-06-30"),
 78    Doc("fin-001", "Car mileage expense",
 79        "Employees who use a personal vehicle for business driving are reimbursed at the standard mileage "
 80        "rate. Log each trip with date, distance and purpose, and submit the expense within 30 days.",
 81        "Finance", "2026-01-05"),
 82    Doc("fin-002", "Travel policy (2026)",
 83        "Book flights and hotels through the travel portal. Economy class is required for flights under "
 84        "six hours. Meals are covered by a daily per-diem of 75 dollars.",
 85        "Finance", "2026-01-01"),
 86    Doc("fin-003", "Travel policy (2023, superseded)",
 87        "Book flights through the travel agency by phone. Business class is allowed for flights over "
 88        "four hours. Meals are covered by a daily per-diem of 60 dollars.",
 89        "Finance", "2023-01-01"),
 90    Doc("fin-004", "Submitting expense receipts",
 91        "Upload receipts to the expense tool within 30 days. Receipts are required for any expense "
 92        "over 25 dollars. Reimbursements are paid with the next payroll run.",
 93        "Finance", "2025-10-12"),
 94    Doc("fin-005", "Vendor invoice reconciliation",
 95        "At quarter end, match each vendor invoice to its payment record. List every mismatch with the "
 96        "invoice number, amount and vendor, and send the reconciliation to the controller.",
 97        "Finance", "2026-03-31", frozenset({"finance"})),
 98    Doc("fin-006", "Q3 revenue forecast",
 99        "Q3 revenue is forecast at 41 million dollars, up 12 percent, driven by enterprise renewals.",
100        "Finance", "2026-07-01", frozenset({"finance", "exec"})),
101    Doc("hr-001", "Paid time off",
102        "Full-time employees accrue 20 days of PTO per year. Request vacation in the HR portal at least "
103        "two weeks ahead. Unused PTO up to 5 days rolls over.",
104        "HR", "2026-01-01"),
105    Doc("hr-002", "Sick leave",
106        "Employees receive 10 paid sick days per year. No manager approval is needed for sick leave, "
107        "but notify your team as early as possible.",
108        "HR", "2025-08-15"),
109    Doc("hr-003", "New hire onboarding",
110        "On your first day, collect your laptop from IT, complete security training, and set up two-factor "
111        "authentication. Your manager will schedule orientation sessions for week one.",
112        "HR", "2026-02-01"),
113    Doc("hr-004", "Salary bands and bonus",
114        "Salary bands are reviewed every April. The annual bonus target is 10 percent of base pay, "
115        "paid in March based on company and individual performance.",
116        "HR", "2026-04-01", frozenset({"hr", "exec"})),
117    Doc("hr-005", "Parental leave",
118        "Primary caregivers receive 16 weeks of paid parental leave; secondary caregivers receive 6 weeks. "
119        "Leave can start up to two weeks before the expected birth or adoption date.",
120        "HR", "2025-05-20"),
121]
122
123DOCS_BY_ID: dict[str, Doc] = {d.id: d for d in DOCS}
124
125
126# (query, set of relevant doc ids). The retrieval "golden set".
127LABELED_QUERIES: list[tuple[str, set[str]]] = [
128    ("I forgot my password and I'm locked out", {"it-001"}),
129    ("how long must a password be", {"it-002"}),
130    ("what does ERR-4012 mean", {"it-004"}),
131    ("ERR-4013 on my laptop", {"it-005"}),
132    ("automobile reimbursement for business driving", {"fin-001"}),
133    ("how many vacation days do I get", {"hr-001"}),
134    ("connect to internal systems from home", {"it-003"}),
135    ("suspicious email with a link", {"it-007"}),
136    ("per-diem for meals when traveling", {"fin-002"}),
137    ("enroll in MFA", {"it-006"}),
138    ("first day checklist for a new hire", {"hr-003"}),
139    ("match vendor bills to payments", {"fin-005"}),
140]