{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "w7-00",
   "metadata": {},
   "source": [
    "# Week 7 Lab — From Role Handoffs to a Research Team\n",
    "\n",
    "**Goal:** Run each specialist yourself before handing the same work to a planner and executor.\n",
    "\n",
    "**Source:** Andrew Ng, *Agentic AI*, Module 5 research-agent project (`GL-M5` / `C1M5_Assignment`). The original tasks and workflow\n",
    "are adapted into guided steps and a final in-class task.\n",
    "\n",
    "| Minutes | Activity | Check before continuing |\n",
    "|---|---|---|\n",
    "| 0–8 | Setup | Research tools and client ready |\n",
    "| 8–20 | Run the researcher | Actual notes and source URLs |\n",
    "| 20–33 | Run writer and editor separately | Research → draft → edited text handoffs |\n",
    "| 33–45 | Plan and choose a role | Task list and a routing decision |\n",
    "| 45–60 | Automate the handoffs | Executor history and a completed report |\n",
    "| 60–77 | In-class task | Own task plan and checked evidence handoff |\n",
    "| 77–80 | Explain and prepare homework | Trace one source through the workflow |\n",
    "\n",
    "Run one cell at a time. Read the actual output at each checkpoint before moving\n",
    "on. Working example values are supplied; the final task asks you to change them\n",
    "and explain the result. The in-class lab is not collected. Homework is listed at\n",
    "the end. All model calls use aisuite; all helpers are included in this notebook.\n",
    "\n",
    "Search adaptation: Ng's Tavily tool is replaced by Week 3's DuckDuckGo `web_search`. The research workflow and returned source links are retained; the search backend and web-result format differ from the original."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-01",
   "metadata": {},
   "source": [
    "## 1. Setup · 0–8 min\n",
    "\n",
    "Run the installation, API-key cell, and imports. Reuse the provider key from Week 3.\n",
    "The default `MODEL` supports tool calling. Web search reuses Week 3's `web_search` (DuckDuckGo via `ddgs`); no separate search API key is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-02",
   "metadata": {
    "tags": [
     "install"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q \"aisuite[openai,anthropic]==0.1.14\" requests ddgs wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-03",
   "metadata": {
    "tags": [
     "credentials"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "MODEL = \"openai:gpt-4o-mini\"\n",
    "OPENAI_API_KEY = \"PASTE-YOUR-KEY-HERE\"\n",
    "if OPENAI_API_KEY != \"PASTE-YOUR-KEY-HERE\":\n",
    "    os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY\n",
    "\n",
    "provider_key = {\"openai\": \"OPENAI_API_KEY\", \"anthropic\": \"ANTHROPIC_API_KEY\"}.get(MODEL.split(\":\")[0])\n",
    "if provider_key and not os.environ.get(provider_key):\n",
    "    raise ValueError(f\"Set {provider_key} before running the lab.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-04",
   "metadata": {
    "tags": [
     "imports"
    ]
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "from datetime import date\n",
    "from IPython.display import Markdown, display\n",
    "import aisuite\n",
    "\n",
    "client = aisuite.Client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-05",
   "metadata": {},
   "source": [
    "## 2. Run the researcher and inspect its output · 8–20 min\n",
    "\n",
    "### 2.1 Give one role one task\n",
    "\n",
    "The research role uses Ng's arXiv and Wikipedia tools and Week 3's web search. The same aisuite\n",
    "loop from Week 4 is automatic inside this role. Run the provided tools, then read\n",
    "the researcher instruction and the task before calling it.\n",
    "\n",
    "The original project topic is the ensemble Kalman filter for time-series forecasting.\n",
    "Predict which sources the researcher will need to describe its assumptions and use.\n",
    "\n",
    "If arXiv returns HTTP 429, the service is limiting requests. Stop repeated retries,\n",
    "keep the error observation, and retry later. In runs with web search available,\n",
    "use returned web evidence and state any missing paper evidence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-06",
   "metadata": {
    "tags": [
     "search-tools"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Original research tools, inlined for Colab — provided { display-mode: \"form\" }\n",
    "import xml.etree.ElementTree as ET\n",
    "import requests\n",
    "from ddgs import DDGS\n",
    "import wikipedia\n",
    "\n",
    "search_session = requests.Session()\n",
    "search_session.headers.update({\"User-Agent\": \"STML-course-research-lab/1.0\"})\n",
    "\n",
    "def arxiv_search_tool(query: str, max_results: int = 5) -> list[dict]:\n",
    "    \"\"\"Search arXiv for academic papers.\n",
    "\n",
    "    Args:\n",
    "        query: Search keywords for research papers.\n",
    "        max_results: Maximum papers to return (1 to 10).\n",
    "    \"\"\"\n",
    "    try:\n",
    "        response = search_session.get(\n",
    "            \"https://export.arxiv.org/api/query\",\n",
    "            params={\"search_query\": f\"all:{query}\", \"start\": 0,\n",
    "                    \"max_results\": max(1, min(max_results, 10))}, timeout=30)\n",
    "        response.raise_for_status()\n",
    "        root = ET.fromstring(response.content)\n",
    "        ns = {\"a\": \"http://www.w3.org/2005/Atom\"}\n",
    "        results = []\n",
    "        for entry in root.findall(\"a:entry\", ns):\n",
    "            results.append({\n",
    "                \"title\": entry.findtext(\"a:title\", default=\"\", namespaces=ns).strip(),\n",
    "                \"authors\": [a.findtext(\"a:name\", namespaces=ns) for a in entry.findall(\"a:author\", ns)],\n",
    "                \"published\": entry.findtext(\"a:published\", default=\"\", namespaces=ns)[:10],\n",
    "                \"url\": entry.findtext(\"a:id\", namespaces=ns),\n",
    "                \"summary\": entry.findtext(\"a:summary\", default=\"\", namespaces=ns).strip(),\n",
    "                \"link_pdf\": next((x.get(\"href\") for x in entry.findall(\"a:link\", ns)\n",
    "                                  if x.get(\"title\") == \"pdf\"), None),\n",
    "            })\n",
    "        return results\n",
    "    except (requests.RequestException, ET.ParseError) as exc:\n",
    "        return [{\"error\": str(exc)}]\n",
    "\n",
    "def web_search(query: str):\n",
    "    \"\"\"Search DuckDuckGo for current information or source links.\n",
    "\n",
    "    Args:\n",
    "        query: Search keywords, for example \"Gyeongbokgung Palace address\".\n",
    "    \"\"\"\n",
    "    try:\n",
    "        results = DDGS(timeout=10).text(query, backend=\"duckduckgo\", max_results=3)\n",
    "        summaries = []\n",
    "        for item in results:\n",
    "            summaries.append(f\"{item['title']}\\nURL: {item['href']}\\n{item['body']}\")\n",
    "        return \"\\n\\n\".join(summaries)\n",
    "    except Exception as exc:\n",
    "        return {\"error\": str(exc)}\n",
    "\n",
    "\n",
    "def wikipedia_search_tool(query: str, sentences: int = 5) -> list[dict]:\n",
    "    \"\"\"Search Wikipedia for an introductory summary.\n",
    "\n",
    "    Args:\n",
    "        query: Wikipedia search keywords.\n",
    "        sentences: Number of summary sentences (1 to 10).\n",
    "    \"\"\"\n",
    "    try:\n",
    "        titles = wikipedia.search(query, results=1)\n",
    "        if not titles:\n",
    "            return []\n",
    "        page = wikipedia.page(titles[0], auto_suggest=False)\n",
    "        summary = wikipedia.summary(page.title, sentences=max(1, min(sentences, 10)), auto_suggest=False)\n",
    "        return [{\"title\": page.title, \"summary\": summary, \"url\": page.url}]\n",
    "    except Exception as exc:\n",
    "        return [{\"error\": str(exc)}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-07",
   "metadata": {
    "tags": [
     "trace-helpers"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Read the actual aisuite messages — provided { display-mode: \"form\" }\n",
    "def response_messages(response):\n",
    "    \"\"\"Convert actual intermediate messages into request-compatible dictionaries.\"\"\"\n",
    "    raw = response.choices[0].intermediate_messages\n",
    "    values = json.loads(json.dumps(raw, default=vars))\n",
    "    allowed = {\"role\", \"content\", \"tool_calls\", \"tool_call_id\", \"name\"}\n",
    "    return [{k: v for k, v in message.items() if k in allowed and v is not None}\n",
    "            for message in values]\n",
    "\n",
    "def show_trace(messages):\n",
    "    for index, message in enumerate(messages):\n",
    "        print(f\"\\n[{index}] {message['role']}\")\n",
    "        for call in message.get(\"tool_calls\", []):\n",
    "            fn = call[\"function\"]\n",
    "            print(\"ACTION\", call[\"id\"], fn[\"name\"], fn[\"arguments\"])\n",
    "        if message.get(\"role\") == \"tool\":\n",
    "            print(\"OBSERVATION\", message.get(\"tool_call_id\"), message.get(\"name\"))\n",
    "        if message.get(\"content\"):\n",
    "            print(message[\"content\"][:1800])  # display only; the full history is retained\n",
    "\n",
    "def final_text(response):\n",
    "    \"\"\"Do not mistake a turn limit with outstanding tool calls for a final report.\"\"\"\n",
    "    message = response.choices[0].message\n",
    "    if getattr(message, \"tool_calls\", None):\n",
    "        raise RuntimeError(\"Tool-turn limit reached before a final answer. Inspect the trace or increase the limit.\")\n",
    "    if not message.content or not message.content.strip():\n",
    "        raise RuntimeError(\"The model returned no final text.\")\n",
    "    return message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-08",
   "metadata": {
    "tags": [
     "researcher"
    ]
   },
   "outputs": [],
   "source": [
    "research_traces = []\n",
    "\n",
    "def research_agent(task, model=MODEL, return_messages=False):\n",
    "    prompt = f\"\"\"You are a research assistant with arxiv_search_tool for papers,\n",
    "web_search for web search, and wikipedia_search_tool for background.\n",
    "Today is {date.today()}. Gather evidence for the task and include returned URLs.\n",
    "Distinguish paper abstracts from full-text evidence. State missing evidence or\n",
    "service failures; do not invent citations.\n",
    "Task: {task}\"\"\"\n",
    "    response = client.chat.completions.create(\n",
    "        model=model, messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "        tools=[arxiv_search_tool, web_search, wikipedia_search_tool],\n",
    "        tool_choice=\"auto\", max_turns=6,\n",
    "    )\n",
    "    messages = response_messages(response)\n",
    "    research_traces.append({\"task\": task, \"messages\": messages})\n",
    "    text = final_text(response)\n",
    "    return (text, messages) if return_messages else text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-09",
   "metadata": {
    "tags": [
     "manual-research"
    ]
   },
   "outputs": [],
   "source": [
    "TOPIC = \"The ensemble Kalman filter for time series forecasting\"\n",
    "research_task = f\"Research {TOPIC}. Find relevant papers and explain the basic method, assumptions, and limitations. Include returned source URLs.\"\n",
    "research_notes, manual_research_messages = research_agent(research_task, return_messages=True)\n",
    "display(Markdown(research_notes))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-10",
   "metadata": {},
   "source": [
    "### 2.2 Inspect a source before forwarding it\n",
    "\n",
    "Find one useful result in the trace. Check that its URL and finding are present\n",
    "in the research notes. If a tool failed or returned only an abstract, keep that\n",
    "limitation visible in the writing task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-11",
   "metadata": {
    "tags": [
     "research-check"
    ]
   },
   "outputs": [],
   "source": [
    "show_trace(manual_research_messages)\n",
    "manual_source_notes = {\"source_url\": \"\", \"finding\": \"\", \"evidence_limit\": \"\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-12",
   "metadata": {},
   "source": [
    "## 3. Run the writer and editor separately · 20–33 min\n",
    "\n",
    "### 3.1 Pass research notes to the writer\n",
    "\n",
    "The writer does not inherit another function's variables. The next cell explicitly\n",
    "includes `research_notes` in its task string. Read that string before the call.\n",
    "Predict which source and limitation should appear in the draft."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-13",
   "metadata": {
    "tags": [
     "writer"
    ]
   },
   "outputs": [],
   "source": [
    "WRITER_SYSTEM = \"\"\"You are a writing agent specialized in well-structured academic\n",
    "or technical content. Use the supplied research context and preserve its source\n",
    "links. Do not invent findings or references. State limits of the available evidence.\"\"\"\n",
    "\n",
    "def writer_agent(task, model=MODEL):\n",
    "    response = client.chat.completions.create(\n",
    "        model=model, messages=[{\"role\": \"system\", \"content\": WRITER_SYSTEM},\n",
    "                               {\"role\": \"user\", \"content\": task}],\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-14",
   "metadata": {
    "tags": [
     "manual-writer-input"
    ]
   },
   "outputs": [],
   "source": [
    "writer_task = f\"\"\"Write a concise report of about 300 words on {TOPIC}.\n",
    "Use these research notes, preserve their source URLs, and state limitations:\n",
    "{research_notes}\"\"\"\n",
    "print(writer_task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-15",
   "metadata": {
    "tags": [
     "manual-writer"
    ]
   },
   "outputs": [],
   "source": [
    "manual_draft = writer_agent(writer_task)\n",
    "display(Markdown(manual_draft))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-16",
   "metadata": {},
   "source": [
    "Check whether the draft uses the source you selected. Mark a claim whose wording\n",
    "needs clarification or qualification before asking the editor.\n",
    "\n",
    "### 3.2 Supply the draft and evidence to the editor\n",
    "\n",
    "The editor needs both the draft to revise and the evidence against which to check\n",
    "its claims. The task below supplies both. Read its input, then inspect the edited\n",
    "report for one specific change and one retained source URL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-17",
   "metadata": {
    "tags": [
     "editor"
    ]
   },
   "outputs": [],
   "source": [
    "EDITOR_SYSTEM = \"\"\"You are an editor agent. Reflect on, critique, or improve the\n",
    "supplied drafts as requested. Check clarity, structure, and whether claims fit\n",
    "the supplied evidence. Preserve source links. If asked for the final document,\n",
    "return the complete revised Markdown report, including limitations and references.\"\"\"\n",
    "\n",
    "def editor_agent(task, model=MODEL):\n",
    "    response = client.chat.completions.create(\n",
    "        model=model, messages=[{\"role\": \"system\", \"content\": EDITOR_SYSTEM},\n",
    "                               {\"role\": \"user\", \"content\": task}],\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-18",
   "metadata": {
    "tags": [
     "manual-editor"
    ]
   },
   "outputs": [],
   "source": [
    "editor_task = f\"\"\"Revise this draft for clarity and supported claims.\n",
    "Return the complete Markdown report with source URLs and limitations.\n",
    "Research notes:\n",
    "{research_notes}\n",
    "Draft:\n",
    "{manual_draft}\"\"\"\n",
    "manual_edited = editor_agent(editor_task)\n",
    "display(Markdown(manual_edited))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-19",
   "metadata": {},
   "outputs": [],
   "source": [
    "manual_handoff_notes = {\"source_in_research\": \"\", \"source_in_draft\": \"\",\n",
    "                        \"source_in_edit\": \"\", \"specific_editor_change\": \"\"}\n",
    "# Fill with actual passages and links from the three outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-20",
   "metadata": {},
   "source": [
    "Checkpoint: identify where `research_notes` entered the writer's input and where\n",
    "`manual_draft` entered the editor's input. Separate role functions can use the same\n",
    "model; their task instructions and supplied context determine their responsibilities.\n",
    "\n",
    "## 4. Generate a plan and choose a role · 33–45 min\n",
    "\n",
    "### 4.1 Inspect the planner's task list\n",
    "\n",
    "The planner returns task strings for research, writing, and editing. The provided\n",
    "parser reads a literal Python list. Before execution, check that each step can be\n",
    "performed by one available role and that the last step returns a complete report."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-21",
   "metadata": {
    "tags": [
     "planner"
    ]
   },
   "outputs": [],
   "source": [
    "import ast\n",
    "\n",
    "PLANNER_INSTRUCTION = \"\"\"You are a planning agent organizing a research workflow.\n",
    "Available agents:\n",
    "- research_agent: search the web, Wikipedia, and arXiv.\n",
    "- writer_agent: draft research summaries and reports.\n",
    "- editor_agent: reflect, critique, and revise drafts.\n",
    "Return ONLY a valid Python list of 3 to 4 task strings. Each step must be atomic\n",
    "and executable by one available agent. Include research, writing, and editing.\n",
    "Do not include creating repositories, installing packages, or unrelated tasks.\n",
    "The final step must produce the complete research report as a Markdown document.\n",
    "Cite sources collected by the researcher and preserve uncertainty.\"\"\"\n",
    "\n",
    "def planner_agent(topic, model=MODEL):\n",
    "    response = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=[{\"role\": \"user\", \"content\": f\"{PLANNER_INSTRUCTION}\\nTopic: {topic}\"}],\n",
    "    )\n",
    "    raw = response.choices[0].message.content\n",
    "    steps = ast.literal_eval(raw.strip())\n",
    "    if not isinstance(steps, list) or not 3 <= len(steps) <= 4 or not all(isinstance(s, str) and s.strip() for s in steps):\n",
    "        raise ValueError(\"Expected 3 to 4 nonempty task strings. Inspect the planner output.\")\n",
    "    return steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-22",
   "metadata": {
    "tags": [
     "plan-run"
    ]
   },
   "outputs": [],
   "source": [
    "plan_steps = planner_agent(TOPIC)\n",
    "for i, step in enumerate(plan_steps, 1):\n",
    "    print(i, step)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-23",
   "metadata": {},
   "source": [
    "### 4.2 Inspect one routing decision\n",
    "\n",
    "Ng's executor asks the model which role should perform a step. That operation is\n",
    "shown as `choose_agent` below so it can be run separately. Predict the role for\n",
    "the first plan step, then compare it with the actual JSON decision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-24",
   "metadata": {
    "tags": [
     "router"
    ]
   },
   "outputs": [],
   "source": [
    "agent_registry = {\"research_agent\": research_agent,\n",
    "                  \"writer_agent\": writer_agent, \"editor_agent\": editor_agent}\n",
    "routing_log = []\n",
    "\n",
    "def clean_json_block(raw):\n",
    "    raw = raw.strip()\n",
    "    if raw.startswith(\"```\"):\n",
    "        raw = re.sub(r\"^```(?:json)?\\s*|\\s*```$\", \"\", raw)\n",
    "    return raw.strip()\n",
    "\n",
    "def choose_agent(step, model=MODEL):\n",
    "    route_prompt = f\"\"\"You are an execution manager for a multi-agent research team.\n",
    "Identify which agent should perform this step and extract its task.\n",
    "Return ONLY a valid JSON object with keys agent and task.\n",
    "agent must be research_agent, writer_agent, or editor_agent.\n",
    "task must preserve the original requirements.\n",
    "Instruction: {step}\"\"\"\n",
    "    response = client.chat.completions.create(\n",
    "        model=model, messages=[{\"role\": \"user\", \"content\": route_prompt}])\n",
    "    decision = json.loads(clean_json_block(response.choices[0].message.content))\n",
    "    if decision.get(\"agent\") not in agent_registry or not isinstance(decision.get(\"task\"), str) or not decision[\"task\"].strip():\n",
    "        raise ValueError(f\"Invalid routing decision: {decision}\")\n",
    "    return decision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-25",
   "metadata": {
    "tags": [
     "route-run"
    ]
   },
   "outputs": [],
   "source": [
    "first_decision = choose_agent(plan_steps[0])\n",
    "display(first_decision)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-26",
   "metadata": {},
   "source": [
    "A routing decision chooses a role and task; it does not execute the role. Check\n",
    "that the selected task retains the original instruction. The execution call and\n",
    "context handoff are assembled next.\n",
    "\n",
    "## 5. Automate the handoffs · 45–60 min\n",
    "\n",
    "### 5.1 Build a role's input from completed history\n",
    "\n",
    "Each history entry stores a step, its role, and the actual output. The helper below\n",
    "places those outputs in the next role's task. This replaces the manual string\n",
    "assembly in Section 3. Inspect its result using the real manual outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-27",
   "metadata": {
    "tags": [
     "context-handoff"
    ]
   },
   "outputs": [],
   "source": [
    "def task_with_context(agent_name, task, history):\n",
    "    context = \"\\n\\n\".join(\n",
    "        f\"Step {j + 1} by {agent}:\\n{output}\"\n",
    "        for j, (_, agent, output) in enumerate(history))\n",
    "    return (f\"You are {agent_name}.\\n\\nContext from completed steps:\\n{context}\"\n",
    "            f\"\\n\\nYour next task:\\n{task}\")\n",
    "\n",
    "manual_history = [(research_task, \"research_agent\", research_notes),\n",
    "                  (\"Write a draft\", \"writer_agent\", manual_draft)]\n",
    "editor_input_from_history = task_with_context(\n",
    "    \"editor_agent\", \"Revise and return the complete report\", manual_history)\n",
    "print(editor_input_from_history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-28",
   "metadata": {},
   "source": [
    "Check that both the research notes and the draft appear in the generated input.\n",
    "If an earlier result is missing from `history`, the helper cannot pass it on.\n",
    "\n",
    "### 5.2 Assemble the executor\n",
    "\n",
    "Read the role-selection call, context construction, role execution, and append\n",
    "operation in order. The completed history grows after each role returns its text.\n",
    "The executor follows the supplied plan and retains its final editing step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-29",
   "metadata": {
    "tags": [
     "executor"
    ]
   },
   "outputs": [],
   "source": [
    "def executor_agent(plan_steps, model=MODEL):\n",
    "    if not isinstance(plan_steps, list) or not 1 <= len(plan_steps) <= 4 or not all(isinstance(s, str) and s.strip() for s in plan_steps):\n",
    "        raise ValueError(\"Supply a nonempty plan of at most four task strings\")\n",
    "    history = []\n",
    "    for i, step in enumerate(plan_steps):\n",
    "        decision = choose_agent(step, model)\n",
    "        agent_name, task = decision[\"agent\"], decision[\"task\"]\n",
    "        enriched_task = task_with_context(agent_name, task, history)\n",
    "        routing_log.append({\"step\": i + 1, \"agent\": agent_name, \"task\": task})\n",
    "        print(f\"Step {i + 1} -> {agent_name}: {task}\")\n",
    "        output = agent_registry[agent_name](enriched_task, model=model)\n",
    "        if not isinstance(output, str) or not output.strip():\n",
    "            raise ValueError(f\"{agent_name} returned no text at step {i + 1}\")\n",
    "        history.append((step, agent_name, output))\n",
    "        display(Markdown(output))\n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-30",
   "metadata": {
    "tags": [
     "run"
    ]
   },
   "outputs": [],
   "source": [
    "research_traces.clear()\n",
    "routing_log.clear()\n",
    "history = executor_agent(plan_steps)\n",
    "final_report = history[-1][2]\n",
    "display(routing_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-31",
   "metadata": {},
   "source": [
    "### 5.3 Compare manual and automatic handoffs\n",
    "\n",
    "The automatic workflow may produce different text from Section 3. Check the same\n",
    "connections: research evidence → writing input → edited report. Open one source\n",
    "and compare the final claim with its supporting passage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-32",
   "metadata": {
    "tags": [
     "automatic-check"
    ]
   },
   "outputs": [],
   "source": [
    "for trace in research_traces:\n",
    "    show_trace(trace[\"messages\"])\n",
    "automatic_handoff_notes = {\"planned_role_vs_actual\": \"\", \"source_url\": \"\",\n",
    "                           \"final_claim\": \"\", \"supporting_passage\": \"\", \"editor_change\": \"\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-33",
   "metadata": {},
   "source": [
    "## 6. In-class task: write a research plan yourself · 60–77 min\n",
    "\n",
    "Design a short plan for a question in your field, then run the existing executor.\n",
    "\n",
    "1. Replace `MY_QUESTION` and write 3–4 task strings in `MY_PLAN`.\n",
    "2. Make one step gather evidence, one write from that evidence, and the final one\n",
    "   return the complete edited report. Specify an output requirement of your own.\n",
    "3. Before execution, predict the role for each step. Compare the actual routing.\n",
    "4. Check one source across the research output, draft, and final report, then verify\n",
    "   one final claim against the source. Record a gap if the evidence is unavailable.\n",
    "\n",
    "This task changes the plan rather than the role functions. The supplied plan is\n",
    "a runnable example; adapt its question and output requirements yourself. Keep\n",
    "every step within the available search, writing, and editing capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-34",
   "metadata": {
    "tags": [
     "challenge-run"
    ]
   },
   "outputs": [],
   "source": [
    "# EDIT: replace the question and plan; keep the final complete-report requirement.\n",
    "MY_QUESTION = \"When do changing dynamics limit ensemble Kalman filter forecasts?\"\n",
    "MY_PLAN = [\n",
    "    f\"Research {MY_QUESTION} using academic and web sources. Return findings, limitations, and source URLs.\",\n",
    "    f\"Write a 300-word report answering {MY_QUESTION}. Use the research notes and include a table of assumptions and possible failures.\",\n",
    "    \"Edit the draft against the supplied evidence. Return the complete Markdown report with the table, limitations, and references.\",\n",
    "]\n",
    "MY_EXPECTED_ROLES = []  # write your prediction before execution\n",
    "challenge_route_start = len(routing_log)\n",
    "challenge_history = executor_agent(MY_PLAN)\n",
    "challenge_report = challenge_history[-1][2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "w7-35",
   "metadata": {
    "tags": [
     "challenge-review"
    ]
   },
   "outputs": [],
   "source": [
    "display(routing_log[challenge_route_start:])\n",
    "my_handoff_review = {\"prediction_vs_routing\": \"\", \"requested_output_present\": \"\",\n",
    "                     \"source_in_research\": \"\", \"source_in_draft\": \"\", \"source_in_report\": \"\",\n",
    "                     \"verified_claim\": \"\", \"supporting_passage\": \"\", \"remaining_gap\": \"\"}\n",
    "# Complete using this workflow's actual outputs and the source you opened."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w7-36",
   "metadata": {},
   "source": [
    "## 7. Explain and prepare homework · 77–80 min\n",
    "\n",
    "Explain how a plan step becomes a role call and where prior research enters the\n",
    "writer's input. Show one source that survived the handoffs, or identify where it\n",
    "was lost. Multiple roles do not by themselves establish report quality.\n",
    "\n",
    "**Homework:** [W7_hw_own_topic.ipynb](W7_hw_own_topic.ipynb). Develop your own topic\n",
    "and checks. Keep intermediate outputs so that failures can be traced to a stage.\n",
    "After the midterm, the same staged workflow motivates LangChain retrieval and\n",
    "LangGraph state management. This lab implements plan-and-execute, not ToT tree search."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "W7_lab_research_agent.ipynb"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "stml": {
   "adaptation": "Ng tasks; guided practice and independent 80-minute classroom task; Week 3 web search",
   "core_minutes": 80,
   "week": 7
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
