{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CS 839: Making Tool Call Work\n",
    "\n",
    "**A quick exercise on tools and tool debugging**\n",
    "\n",
    "You are given a Python tool (that already works). Your goal is to make its interface to the model work reliably.\n",
    "\n",
    "By the end, you should be able to understand four stages of tool use: tool selection, argument construction, local execution, and result interpretation.\n",
    "\n",
    "Our scientific task: Find the experiment with the highest score among runs below 50 °C using material A. Give the run ID and supporting value. We will use tool-calling agents to do this. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Steps\n",
    "\n",
    "- Inspect the toy data and tool implementation.\n",
    "- Run the intentionally weak interface.\n",
    "- Repair the interface.\n",
    "- If we have time, try an adversarial prompt from another group.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Run the next two cells. When prompted, paste your team's API key. **Good practice: never paste a key into a normal code cell or Git repo.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"openai>=1.0\" pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import json\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "from openai import OpenAI\n",
    "\n",
    "MODEL = os.getenv(\"CS839_MODEL\", \"gpt-5.6-luna\")\n",
    "\n",
    "if not os.getenv(\"OPENAI_API_KEY\"):\n",
    "    entered_key = getpass.getpass(\"Paste your team's API key (hidden): \").strip()\n",
    "    if entered_key:\n",
    "        os.environ[\"OPENAI_API_KEY\"] = entered_key\n",
    "\n",
    "client = OpenAI() if os.getenv(\"OPENAI_API_KEY\") else None\n",
    "print(f\"Model: {MODEL}\")\n",
    "print(\"API client ready.\" if client else \"No key supplied. Local cells will still work.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Toy experiment data\n",
    "\n",
    "These are made-up (sythetic) runs. Larger scores are better. The table is small enough that you can verify every answer yourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENTS = [\n",
    "    {\"run_id\": \"A-01\", \"material\": \"A\", \"temperature_c\": 35.0, \"score\": 71.0},\n",
    "    {\"run_id\": \"A-02\", \"material\": \"A\", \"temperature_c\": 44.0, \"score\": 82.0},\n",
    "    {\"run_id\": \"A-03\", \"material\": \"A\", \"temperature_c\": 49.0, \"score\": 91.0},\n",
    "    {\"run_id\": \"A-04\", \"material\": \"A\", \"temperature_c\": 50.0, \"score\": 96.0},\n",
    "    {\"run_id\": \"A-05\", \"material\": \"A\", \"temperature_c\": 52.0, \"score\": 99.0},\n",
    "    {\"run_id\": \"B-01\", \"material\": \"B\", \"temperature_c\": 36.0, \"score\": 76.0},\n",
    "    {\"run_id\": \"B-02\", \"material\": \"B\", \"temperature_c\": 45.0, \"score\": 89.0},\n",
    "    {\"run_id\": \"B-03\", \"material\": \"B\", \"temperature_c\": 49.5, \"score\": 93.0},\n",
    "    {\"run_id\": \"B-04\", \"material\": \"B\", \"temperature_c\": 51.0, \"score\": 97.0},\n",
    "    {\"run_id\": \"C-01\", \"material\": \"C\", \"temperature_c\": 40.0, \"score\": 79.0},\n",
    "    {\"run_id\": \"C-02\", \"material\": \"C\", \"temperature_c\": 48.0, \"score\": 87.0},\n",
    "    {\"run_id\": \"C-03\", \"material\": \"C\", \"temperature_c\": 53.0, \"score\": 98.0},\n",
    "]\n",
    "\n",
    "experiments_df = pd.DataFrame(EXPERIMENTS)\n",
    "experiments_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The local tool\n",
    "\n",
    "This is deterministic Python code. As we have discussed in class, the model cannot execute it---the model can only request that our program (harness) call it. The temperature bound we are looking for is **strictly less than** the supplied Celsius value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_runs(*, material=None, temperature=None, max_temperature_c=None):\n",
    "    \"\"\"Return matching rows ordered from highest to lowest score.\"\"\"\n",
    "    threshold = max_temperature_c if max_temperature_c is not None else temperature\n",
    "\n",
    "    if threshold is None:\n",
    "        return {\"ok\": False, \"error\": \"A maximum temperature is required.\"}\n",
    "    if material is not None and material not in {\"A\", \"B\", \"C\"}:\n",
    "        return {\"ok\": False, \"error\": f\"Unknown material: {material}\"}\n",
    "\n",
    "    rows = experiments_df[experiments_df[\"temperature_c\"] < float(threshold)]\n",
    "    if material is not None:\n",
    "        rows = rows[rows[\"material\"] == material]\n",
    "    rows = rows.sort_values(\"score\", ascending=False)\n",
    "\n",
    "    return {\n",
    "        \"ok\": True,\n",
    "        \"filter\": {\n",
    "            \"material\": material,\n",
    "            \"temperature_rule\": f\"temperature_c < {float(threshold)}\",\n",
    "        },\n",
    "        \"matches\": rows.to_dict(orient=\"records\"),\n",
    "    }\n",
    "\n",
    "# Verify the implementation without spending API credit.\n",
    "local_check = find_runs(material=\"A\", max_temperature_c=50)\n",
    "assert local_check[\"matches\"][0][\"run_id\"] == \"A-03\"\n",
    "local_check"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. An intentionally flawed tool interface\n",
    "\n",
    "The function works, but this interface leaves important questions unanswered. Read it carefully before running anything."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BAD_TOOL = {\n",
    "    \"type\": \"function\",\n",
    "    \"name\": \"search_runs\",\n",
    "    \"description\": \"Searches the experiment data.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"material\": {\"type\": \"string\"},\n",
    "            \"temperature\": {\"type\": \"number\"},\n",
    "        },\n",
    "        \"required\": [\"temperature\"],\n",
    "        \"additionalProperties\": True,\n",
    "    },\n",
    "    \"strict\": False,\n",
    "}\n",
    "\n",
    "BAD_TOOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Small agent loop\n",
    "\n",
    "Follow the printed trace: user request → model tool call → local Python result → model answer. You do not need to modify this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_local_tool(name, arguments):\n",
    "    if name not in {\"search_runs\", \"find_experimental_runs\"}:\n",
    "        return {\"ok\": False, \"error\": f\"Unknown tool: {name}\"}\n",
    "    return find_runs(\n",
    "        material=arguments.get(\"material\"),\n",
    "        temperature=arguments.get(\"temperature\"),\n",
    "        max_temperature_c=arguments.get(\"max_temperature_c\"),\n",
    "    )\n",
    "\n",
    "\n",
    "def run_agent(prompt, tool_definition, max_steps=3):\n",
    "    if client is None:\n",
    "        raise RuntimeError(\"No API key is configured. Re-run the setup cell.\")\n",
    "\n",
    "    input_items = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    trace = {\"prompt\": prompt, \"calls\": [], \"final_answer\": None}\n",
    "    total_input_tokens = 0\n",
    "    total_output_tokens = 0\n",
    "    print(f\"USER: {prompt}\\n\")\n",
    "\n",
    "    for step in range(1, max_steps + 1):\n",
    "        response = client.responses.create(\n",
    "            model=MODEL,\n",
    "            instructions=(\n",
    "                \"You are a careful assistant analyzing a tiny fictional experiment table. \"\n",
    "                \"Use the available tool for claims about the data. Never invent records. \"\n",
    "                \"Give the winning run ID and score when the request is answerable.\"\n",
    "            ),\n",
    "            input=input_items,\n",
    "            tools=[tool_definition],\n",
    "            tool_choice=\"auto\",\n",
    "        )\n",
    "\n",
    "        if response.usage:\n",
    "            total_input_tokens += response.usage.input_tokens\n",
    "            total_output_tokens += response.usage.output_tokens\n",
    "\n",
    "        input_items += response.output\n",
    "        function_calls = [item for item in response.output if item.type == \"function_call\"]\n",
    "\n",
    "        if not function_calls:\n",
    "            answer = response.output_text or \"(No textual answer returned.)\"\n",
    "            trace[\"final_answer\"] = answer\n",
    "            trace[\"usage\"] = {\n",
    "                \"input_tokens\": total_input_tokens,\n",
    "                \"output_tokens\": total_output_tokens,\n",
    "            }\n",
    "            print(f\"ASSISTANT: {answer}\")\n",
    "            print(f\"\\nTOKENS: {trace['usage']}\")\n",
    "            return trace\n",
    "\n",
    "        for call in function_calls:\n",
    "            try:\n",
    "                arguments = json.loads(call.arguments)\n",
    "            except json.JSONDecodeError as exc:\n",
    "                arguments = {}\n",
    "                result = {\"ok\": False, \"error\": f\"Invalid JSON arguments: {exc}\"}\n",
    "            else:\n",
    "                result = execute_local_tool(call.name, arguments)\n",
    "\n",
    "            trace[\"calls\"].append(\n",
    "                {\"step\": step, \"name\": call.name, \"arguments\": arguments, \"result\": result}\n",
    "            )\n",
    "            print(f\"TOOL REQUEST: {call.name}({json.dumps(arguments, sort_keys=True)})\")\n",
    "            print(f\"TOOL RESULT:  {json.dumps(result, sort_keys=True)}\\n\")\n",
    "            input_items.append({\n",
    "                \"type\": \"function_call_output\",\n",
    "                \"call_id\": call.call_id,\n",
    "                \"output\": json.dumps(result),\n",
    "            })\n",
    "\n",
    "    trace[\"final_answer\"] = \"Stopped after the maximum number of tool steps.\"\n",
    "    trace[\"usage\"] = {\n",
    "        \"input_tokens\": total_input_tokens,\n",
    "        \"output_tokens\": total_output_tokens,\n",
    "    }\n",
    "    print(trace[\"final_answer\"])\n",
    "    return trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run the weak interface\n",
    "\n",
    "Run the clear prompt and then the edge case prompts. Inspect the tool arguments, not only the final written answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROMPTS = {\n",
    "    \"clear\": \"Which material A run has the highest score among runs below 50 degrees Celsius?\",\n",
    "    \"kelvin\": \"Which material A run has the highest score among runs below 323 K?\",\n",
    "    \"missing_material\": \"Which run has the highest score below 50 degrees Celsius?\",\n",
    "    \"ambiguous_unit\": \"Find the best material A run below 320.\",\n",
    "}\n",
    "\n",
    "baseline_clear = run_agent(PROMPTS[\"clear\"], BAD_TOOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose edge case prompts by changing this key.\n",
    "chosen_edge_case = \"kelvin\"\n",
    "baseline_edge = run_agent(PROMPTS[chosen_edge_case], BAD_TOOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Diagnose before editing\n",
    "\n",
    "Answer the following: Did the model call the tool? What exact arguments did it generate? Were units and the boundary represented correctly? If something failed, classify it as **selection**, **arguments**, **execution**, or **interpretation**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Repair the interface\n",
    "\n",
    "Complete the TODO descriptions below. You may also change the name, enum, required fields, or other schema details. Decide which values are valid, what unit is accepted, whether the maximum is inclusive, and when the agent should ask a clarification question.\n",
    "\n",
    "Strict schemas require every property to appear in required and set additionalProperties to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STUDENT_TOOL = {\n",
    "    \"type\": \"function\",\n",
    "    \"name\": \"find_experimental_runs\",\n",
    "    \"description\": (\n",
    "        \"TODO: Explain exactly when to use this tool and when to ask the user \"\n",
    "        \"a clarification question instead.\"\n",
    "    ),\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"material\": {\n",
    "                \"type\": \"string\",\n",
    "                \"enum\": [\"A\", \"B\", \"C\"],\n",
    "                \"description\": \"TODO: Describe the material argument.\",\n",
    "            },\n",
    "            \"max_temperature_c\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"TODO: State the unit and whether the boundary is inclusive.\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"material\", \"max_temperature_c\"],\n",
    "        \"additionalProperties\": False,\n",
    "    },\n",
    "    \"strict\": True,\n",
    "}\n",
    "\n",
    "STUDENT_TOOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Test the repaired interface\n",
    "\n",
    "Run the same edge case again, then replace the adversarial prompt with one from another group. A good outcome is not always a tool call; sometimes the correct behavior is to ask a question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repaired_edge = run_agent(PROMPTS[chosen_edge_case], STUDENT_TOOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace this with a prompt supplied by another group.\n",
    "adversarial_prompt = \"Find the best material A run below 320.\"\n",
    "adversarial_trace = run_agent(adversarial_prompt, STUDENT_TOOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Finish\n",
    "\n",
    "Produce the following answer:\n",
    "\n",
    "1. **Failure:** original prompt and bad arguments or behavior.\n",
    "2. **Diagnosis:** selection, arguments, execution, or interpretation.\n",
    "3. **Repair:** one interface change that mattered.\n",
    "4. **Result:** what happened afterward.\n",
    "5. **Open question:** one ambiguity the interface still cannot represent.\n",
    "\n",
    "---\n",
    "\n",
    "[OpenAI function-calling guide](https://developers.openai.com/api/docs/guides/function-calling)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "CS839 Make the Tool Call Work.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}