{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Lab 2 Answer Key*\n",
    "\n",
    "# Multi-Agent Systems & Orchestration\n",
    "\n",
    "**Objective:** decide single-agent versus orchestrator-worker for a discovery system, and write the delegation brief that keeps the split safe.\n",
    "\n",
    "This is a design lab: the deliverables are a justified topology choice and a self-contained delegation brief, with one small cost model at the end.\n",
    "\n",
    "**Scenario.** A litigation team runs discovery over 200,000 documents: emails, contracts, scanned PDFs. For a given matter they need every document responsive to a set of legal criteria, each flagged with the criterion it matches and a one-line justification a paralegal can verify. A privileged document that is accidentally produced is a serious, sometimes sanctionable, event. Turnaround is measured in hours, not seconds. The corpus does not fit in a single context window.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Run the bar: parallelism and isolation\n",
    "\n",
    "Multi-agent needs genuine parallelism or context isolation, and isolation is the stronger reason. Write one sentence for each half before you choose anything, and note that '200,000 documents exceed one context window' is a real answer, not a hand-wave.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Parallelism:** yes. Reviewing batch 1 against the criteria needs nothing from the review of batch 2; the subtasks are genuinely independent, and turnaround measured in hours rewards running them concurrently.\n",
    "\n",
    "**Context isolation:** yes, and this is the load-bearing half. 200,000 documents at even 1,500 tokens each is roughly 300M tokens; no context window carries that. Each worker must carry only its own batch, and the parent must carry only the distilled findings, never the corpus.\n",
    "\n",
    "Both halves of the bar pass. This is the rare scenario where multi-agent is the correct answer rather than the impressive one, and it is worth noticing that the scenario had to work hard to earn it: a corpus bigger than any window, independent units of work, and an hours-scale deadline.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Reject the naive single agent explicitly\n",
    "\n",
    "Say the sentence: 'One agent cannot do this because ___.' If your reason is only 'it would be slow,' you have not found the isolation seam; look again.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**One agent cannot do this because the corpus exceeds any context window, so a single agent reading documents sequentially would fill its window with document text long before finishing, and would have to discard earlier documents to admit later ones.** At that point its judgments are made against a sliding fragment of the corpus, and context editing or compaction would be discarding exactly the material it is being asked to judge.\n",
    "\n",
    "Slowness is real but it is the weak objection: a slow correct system is acceptable at an hours-scale deadline. The disqualifier is the isolation seam. The work must be split so that each reasoning loop carries one batch and the parent carries only findings; that is Lab 2's forty-files principle applied at discovery scale.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Choose the topology and justify it\n",
    "\n",
    "Orchestrator-worker (a manager fans out document batches, workers return responsive hits) versus a flat parallel fan-out. Which failure mode from the coordination table are you most exposed to, and how does your topology contain it?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Topology: orchestrator-worker.** The manager partitions the corpus into non-overlapping batches, dispatches one self-contained brief per batch, and a single synthesize step owns every write to the responsive-documents log. A flat fan-out with no manager has no owner for partitioning (inviting duplicated work at the batch seams) and no single writer for the results log (inviting conflicting writes). Hierarchical is unjustified: no worker's batch is itself large enough to decompose, so added depth would be architecture theatre.\n",
    "\n",
    "| Failure mode | Exposure here | Contained by |\n",
    "|---|---|---|\n",
    "| **Lost constraint** | **Highest**: the privilege rule lives with the manager, and a worker that never receives it flags privileged documents as responsive. | A versioned brief template with the privilege rule baked in (Part 4), never re-typed per dispatch. |\n",
    "| Duplicated work | Batches could overlap at partition seams. | The manager assigns explicit, non-overlapping document ID ranges. |\n",
    "| Conflicting writes | Many workers, one results log. | Workers return findings; only the synthesize step writes. |\n",
    "| Cost blowup | Every worker needs the same criteria and instructions. | Cache the shared prefix (Part 6). |\n",
    "| Infinite delegation | Low, but cheap to prevent. | Workers cannot spawn workers; depth is capped at one. |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Write the worker's delegation brief\n",
    "\n",
    "A worker sees one batch and nothing else; it inherits nothing from the manager's conversation. Enumerate exactly what must be in the prompt: the legal criteria, the definition of responsive, the output contract, and, critically, the privilege rule.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The brief is the interface contract. It is a VERSIONED TEMPLATE: the\n",
    "# privilege rule is baked into the template, never re-typed per dispatch,\n",
    "# so it cannot be the constraint the manager forgets.\n",
    "CRITERIA = [\n",
    "    \"C1: communications concerning the Delta contract negotiation, 2019-2021\",\n",
    "    \"C2: documents discussing warranty claims for the K-series product line\",\n",
    "    \"C3: internal analyses of Southeast-region market share\",\n",
    "]\n",
    "\n",
    "RESPONSIVE_DEF = (\n",
    "    \"A document is responsive if its content substantively addresses at least \"\n",
    "    \"one criterion. A passing keyword mention without substantive discussion \"\n",
    "    \"is NOT responsive. When in doubt, mark responsive and say why: \"\n",
    "    \"over-inclusion is reviewed by a human; silent omission is not.\"\n",
    ")\n",
    "\n",
    "PRIVILEGE_RULE = (\n",
    "    \"PRIVILEGE (always applies): if a document is to, from, or copying counsel, \"\n",
    "    \"or contains legal advice, set privilege_suspected=true and do NOT quote \"\n",
    "    \"its contents in the justification. Your flag routes the document to \"\n",
    "    \"attorney review; it is not a privilege determination.\"\n",
    ")\n",
    "\n",
    "OUTPUT_CONTRACT = (\n",
    "    \"Return one JSON object per responsive document, and nothing for \"\n",
    "    \"non-responsive documents: {\\\"doc_id\\\": str, \\\"criterion\\\": \\\"C1|C2|C3\\\", \"\n",
    "    \"\\\"justification\\\": one line a paralegal can verify, \"\n",
    "    \"\\\"privilege_suspected\\\": bool}. Do not echo document text back.\"\n",
    ")\n",
    "\n",
    "\n",
    "def build_brief(batch_id: str, doc_id_range: str) -> str:\n",
    "    \"\"\"Everything the worker needs, because it starts with an empty context.\"\"\"\n",
    "    return (\n",
    "        f\"TASK: review discovery batch {batch_id} (documents {doc_id_range}, \"\n",
    "        \"attached below) against the criteria.\\n\\n\"\n",
    "        \"CRITERIA:\\n\" + \"\\n\".join(f\"- {c}\" for c in CRITERIA) + \"\\n\\n\"\n",
    "        f\"RESPONSIVE MEANS: {RESPONSIVE_DEF}\\n\\n\"\n",
    "        f\"{PRIVILEGE_RULE}\\n\\n\"\n",
    "        f\"RETURN EXACTLY: {OUTPUT_CONTRACT}\\n\"\n",
    "    )\n",
    "\n",
    "\n",
    "print(build_brief(\"B-0042\", \"58200-58249\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If the privilege rule is the constraint the manager forgets:** the worker reasons correctly over an incomplete brief, confidently marks privileged documents responsive with well-formed justifications that quote the privileged text, and nothing anywhere throws. Those documents flow toward the production set looking exactly like every legitimate hit. That is the lost-constraint bug with a sanctions hearing attached, which is why the rule lives in the versioned template rather than in anyone's memory, and why the human gate in Part 5 exists even so.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Handle the privileged-document blast radius\n",
    "\n",
    "A model flag is not a privilege determination. Where is the human gate, and does it sit on every produced document or only the ambiguous ones? Tie this back to Lab 1's egress plane.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The gate sits on every produced document, not only the flagged ones.** Production is the irreversible action in this system (the analogue of Lab 1's denial branch), and the blast radius of one missed privileged document is sanctionable. A gate keyed to `privilege_suspected == true` would make the model's recall the firm's last line of defense, and the flag exists to route review, not to discharge it.\n",
    "\n",
    "Egress design:\n",
    "\n",
    "- **Every candidate production document** passes attorney or supervised-paralegal review before production. The model's criterion and justification make that review fast; they do not replace it.\n",
    "- **`privilege_suspected` ranks and routes**: flagged documents go to a senior privilege-review queue first. The flag orders the queue; it never bypasses one. This is Lab 1's confidence-score rule wearing a legal hat.\n",
    "- **Schema validation at egress**: any worker output that fails the output contract (missing doc_id, unknown criterion, quoted text where privilege was suspected) is rejected and the batch re-queued, before a human ever reads it.\n",
    "- **The audit trail**: every finding carries its batch ID, brief version, and model version, so when opposing counsel asks how a document was identified, the answer is reconstructable.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: Defend the cost\n",
    "\n",
    "200,000 documents across N workers. What is the shared prefix, where does the cache breakpoint go, and roughly what does caching save versus each worker reloading the criteria cold? State the number you would put in front of the partner who signs the discovery bill.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The shared prefix is everything identical across workers: the review instructions, the criteria, the definition of responsive, the privilege rule, and the output contract. It renders first in every worker's context with a `cache_control` breakpoint on its last block; each worker's batch text follows it. The first dispatch writes the cache at 1.25x; a steady dispatch stream keeps the 5-minute entry warm, so the remaining workers read the prefix at roughly 0.1x input price.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cost model: Haiku 4.5 workers ($1 in / $5 out per 1M tokens). The task is\n",
    "# well-specified classification against explicit criteria: Haiku territory.\n",
    "HAIKU_IN = 1.00 / 1_000_000\n",
    "HAIKU_OUT = 5.00 / 1_000_000\n",
    "CACHE_READ, CACHE_WRITE = 0.1, 1.25\n",
    "\n",
    "DOCS = 200_000\n",
    "DOC_TOKENS = 1_500                 # average email/contract/PDF page text\n",
    "BATCH = 50\n",
    "WORKERS = DOCS // BATCH            # 4,000 worker dispatches\n",
    "PREFIX = 8_000                     # instructions + criteria + contract (shared)\n",
    "FINDING_TOKENS = 400               # per-worker output (hits only, no echo)\n",
    "\n",
    "# Shared prefix: cold reload vs cached (continuous dispatch keeps it warm).\n",
    "prefix_uncached = WORKERS * PREFIX * HAIKU_IN\n",
    "prefix_cached = (PREFIX * CACHE_WRITE\n",
    "                 + (WORKERS - 1) * PREFIX * CACHE_READ) * HAIKU_IN\n",
    "\n",
    "# The parts caching cannot touch: batch text in, findings out.\n",
    "doc_cost = DOCS * DOC_TOKENS * HAIKU_IN\n",
    "out_cost = WORKERS * FINDING_TOKENS * HAIKU_OUT\n",
    "\n",
    "total = prefix_cached + doc_cost + out_cost\n",
    "print(f\"prefix: ${prefix_uncached:6.2f} uncached -> ${prefix_cached:5.2f} cached \"\n",
    "      f\"({1 - prefix_cached / prefix_uncached:.0%} saved on the prefix)\")\n",
    "print(f\"documents in: ${doc_cost:,.2f}   findings out: ${out_cost:,.2f}\")\n",
    "print(f\"total model spend for the review: ${total:,.2f}\")\n",
    "\n",
    "# The number for the partner: first-pass human review at ~$1/document.\n",
    "print(f\"first-pass human review baseline: ${DOCS * 1.00:,.0f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two honest observations belong in the writeup. First, caching saves about 90% of the *prefix* cost, but on this workload the document text dominates the bill, so the headline number is the total: roughly $350 of model spend to first-pass 200,000 documents, against a human first-pass baseline near $200,000. Second, the model pass does not remove the attorneys; it removes the linear read. The partner's number is the review hours redirected from 200,000 documents to the responsive and privilege-flagged subset, with the gate from Part 5 intact.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
