{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ⚡ Level 1.5 · Session 2 — Powering the Agent\n",
    "### AI Agents track · Google Colab · powered by **Google Gemini (free)**\n",
    "\n",
    "Last time you built an agent. Now you'll **power it up**: more **tools**, sharper **prompts**, and a real **memory**.\n",
    "\n",
    "| Part | You'll do | ✅ Done when |\n",
    "|---|---|---|\n",
    "| **A · Tools** | Wire your own Python tools into the agent loop | The agent calls your tool to reach a goal |\n",
    "| **B · Prompting** | Compare a vague prompt vs a sharp one | The sharp prompt visibly changes the reply |\n",
    "| **C · Memory** | Store facts, retrieve the right one | The agent answers using a remembered fact |\n",
    "\n",
    "> Run each grey cell in order with ▶️ / Shift+Enter. Read the note above it first."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 🛠️ Setup — same engine as Session 1\n",
    "### 🔑 Get your free Gemini key (once)\n",
    "1. Open **https://aistudio.google.com/apikey** → **Create API key** → copy it.\n",
    "2. Run the cells below; paste the key when asked.\n",
    "\n",
    "🔒 An API key is a password — never paste it into a public doc, chat, or GitHub."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip -q install google-genai"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from google import genai\n",
    "from google.genai import types\n",
    "import getpass\n",
    "API_KEY = getpass.getpass('Paste your Gemini API key: ').strip()\n",
    "client = genai.Client(api_key=API_KEY)\n",
    "print('Client ready ✅')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Auto-pick models (names change, so we ask your key what's available)\n",
    "model_names = [m.name for m in client.models.list()]\n",
    "flash = [m for m in model_names if 'flash' in m and 'embedding' not in m]\n",
    "MODEL = flash[0] if flash else 'models/gemini-2.5-flash'\n",
    "embeds = [m for m in model_names if 'embedding' in m]\n",
    "EMBED_MODEL = embeds[0] if embeds else 'models/text-embedding-004'\n",
    "print('Using MODEL =', MODEL)\n",
    "print('Using EMBED_MODEL =', EMBED_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ✅ TEST: the engine works\n",
    "r = client.models.generate_content(model=MODEL, contents='In 8 words, what makes an agent powerful?')\n",
    "print(r.text)\n",
    "assert r.text, 'No reply — check your key/model.'\n",
    "print('\\n✅ Setup works! Now let\\'s power up the agent.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 🅰️ Part A — Tools via function calling\n",
    "A **tool** is just a plain Python function the model can choose to call. Let's give the agent two: a **word_counter** and a mock **get_weather**. Then we wire them into the SAME manual loop as Session 1 — auto-calling OFF so *we* watch every **Act → Observe** step."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def word_counter(text: str) -> str:\n",
    "    \"\"\"Count the number of words in a piece of text.\"\"\"\n",
    "    return str(len(text.split()))\n",
    "\n",
    "def get_weather(city: str) -> str:\n",
    "    \"\"\"Return the (mock) current weather for a city.\"\"\"\n",
    "    fake = {'austin': 'sunny, 95°F', 'seattle': 'rainy, 58°F', 'denver': 'snowy, 30°F'}\n",
    "    return fake.get(city.strip().lower(), 'clear, 72°F')\n",
    "\n",
    "print('Tools ready:', word_counter.__name__, '+', get_weather.__name__)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TOOLS = {'word_counter': word_counter, 'get_weather': get_weather}\n",
    "\n",
    "config = types.GenerateContentConfig(\n",
    "    system_instruction=(\n",
    "        'You are a helpful agent. You cannot count words reliably and you do not know the '\n",
    "        'weather, so you MUST call the word_counter and get_weather tools instead of guessing.'),\n",
    "    tools=[word_counter, get_weather],\n",
    "    automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),\n",
    ")\n",
    "\n",
    "goal = 'How many words are in \\'the quick brown fox jumps\\', and what is the weather in Austin?'\n",
    "history = [types.Content(role='user', parts=[types.Part(text=goal)])]\n",
    "\n",
    "for step in range(6):  # safety cap so it can't loop forever\n",
    "    resp = client.models.generate_content(model=MODEL, contents=history, config=config)\n",
    "    parts = resp.candidates[0].content.parts\n",
    "    calls = [p.function_call for p in parts if p.function_call]\n",
    "    if not calls:\n",
    "        print('\\n🏁 FINAL ANSWER:\\n', resp.text); break\n",
    "    history.append(resp.candidates[0].content)          # Observe/Memory: remember the action\n",
    "    for fc in calls:\n",
    "        print(f'🛠️  Act → {fc.name}({dict(fc.args)})')\n",
    "        result = TOOLS[fc.name](**dict(fc.args))\n",
    "        print(f'👀 Observe → {result}')\n",
    "        history.append(types.Content(role='user',\n",
    "            parts=[types.Part.from_function_response(name=fc.name, response={'result': result})]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🧪 **Test:** you should see `🛠️ Act` and `👀 Observe` lines for BOTH tools, then a correct final answer (5 words · Austin sunny). More tools = a more capable agent."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 🎯 YOUR TURN: add a THIRD tool and give a goal that needs it.\n",
    "def to_uppercase(text: str) -> str:\n",
    "    \"\"\"Convert text to UPPERCASE.\"\"\"\n",
    "    return text.upper()\n",
    "\n",
    "# Add to_uppercase to tools=[...] and TOOLS above, then try a goal like:\n",
    "# 'Shout \\'go team\\' in uppercase and tell me the weather in Denver.'\n",
    "print('Wire to_uppercase into the agent and re-run the loop!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 🅱️ Part B — Prompting: vague vs sharp\n",
    "Same model, same question — but the **system prompt** steers the behavior. A **sharp** prompt gives a **Role + rules + format**. Let's run both on the same task and compare."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "task = 'Give me tips for my first day of high school.'\n",
    "\n",
    "vague = 'You are helpful.'\n",
    "\n",
    "sharp = (\n",
    "    'ROLE: You are a friendly coach for nervous 9th graders.\\n'\n",
    "    'RULES: Be encouraging. No jargon. Keep it practical.\\n'\n",
    "    'FORMAT: Reply as EXACTLY 3 bullet points, each starting with an emoji, max 12 words each.')\n",
    "\n",
    "def ask(system_prompt, task):\n",
    "    cfg = types.GenerateContentConfig(system_instruction=system_prompt)\n",
    "    return client.models.generate_content(model=MODEL, contents=task, config=cfg).text\n",
    "\n",
    "print('🟡 VAGUE PROMPT reply:\\n', ask(vague, task))\n",
    "print('\\n' + '='*50 + '\\n')\n",
    "print('🟢 SHARP PROMPT reply:\\n', ask(sharp, task))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🧪 **Test:** the sharp reply should look visibly different — exactly 3 emoji bullets, short and practical — while the vague one rambles. Same model; the **prompt** changed the behavior."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 🎯 YOUR TURN: improve this weak prompt so the reply is a short, kid-friendly rhyme.\n",
    "weak = 'Write about recycling.'\n",
    "\n",
    "# Rewrite `weak` with a ROLE + RULES + FORMAT, then run:\n",
    "# print(ask(weak, 'Write about recycling.'))\n",
    "print('Sharpen the weak prompt and compare the replies!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 🅲 Part C — Memory with embeddings\n",
    "An agent with **memory** can recall facts it was told earlier. We turn each fact into a list of numbers (an **embedding**), then find the fact whose numbers are *closest* to the question. That's the core of long-term memory (aka **RAG**)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# A tiny fake student profile — our agent's 'memory'\n",
    "facts = [\n",
    "    'The student\\'s name is Priya and she is in 9th grade.',\n",
    "    'Priya\\'s favorite subject is biology.',\n",
    "    'Priya is on the school soccer team and plays goalkeeper.',\n",
    "    'Priya wants to become a marine biologist.',\n",
    "]\n",
    "\n",
    "def embed(texts):\n",
    "    e = client.models.embed_content(model=EMBED_MODEL, contents=texts)\n",
    "    return [np.array(x.values) for x in e.embeddings]\n",
    "\n",
    "fact_vecs = embed(facts)\n",
    "print('Embedded', len(fact_vecs), 'facts — each is', len(fact_vecs[0]), 'numbers long.')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def cosine(a, b):\n",
    "    \"\"\"Similarity of two vectors: 1.0 = identical, 0 = unrelated.\"\"\"\n",
    "    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))\n",
    "\n",
    "question = 'What position does Priya play in sports?'\n",
    "q_vec = embed([question])[0]\n",
    "\n",
    "scores = [cosine(q_vec, fv) for fv in fact_vecs]\n",
    "best = int(np.argmax(scores))\n",
    "retrieved = facts[best]\n",
    "print('🔎 Retrieved fact:', retrieved)\n",
    "print('   (score:', round(scores[best], 3), ')')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Answer WITHOUT memory vs WITH the retrieved fact as context\n",
    "no_mem = client.models.generate_content(model=MODEL, contents=question).text\n",
    "\n",
    "prompt = f'Context: {retrieved}\\n\\nQuestion: {question}\\nAnswer using only the context.'\n",
    "with_mem = client.models.generate_content(model=MODEL, contents=prompt).text\n",
    "\n",
    "print('🟡 WITHOUT memory:\\n', no_mem)\n",
    "print('\\n🟢 WITH memory:\\n', with_mem)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🧪 **Test:** the retrieved fact should be the soccer one, and the **WITH memory** answer says *goalkeeper* — while **WITHOUT memory** the model can only guess. That's retrieval-powered memory. 🧠"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 🎯 YOUR TURN: add your OWN facts, then ask a question about them.\n",
    "# facts.append('Priya has a pet dog named Comet.')\n",
    "# fact_vecs = embed(facts)   # re-embed after changing facts\n",
    "# ...then re-run the retrieval cell with a new question like 'What pet does Priya have?'\n",
    "print('Add a fact, re-embed, and ask about it!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 🎉 You powered up your agent\n",
    "**Tools** give it more hands, **prompts** steer its brain, and **memory** lets it remember. Same loop as Session 1 — now much more capable.\n",
    "\n",
    "**Next session — Agents that Plan:** breaking a big goal into steps, chaining tools, and letting the agent decide the order.\n",
    "\n",
    "_AI Agents track · LearnAI Level 1.5 · part of MillionRoots._"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "L1.5 Session 2 — Powering the Agent",
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}