{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5a3dca7d-514b-4e75-a198-74851a73aed3",
   "metadata": {},
   "source": [
    "## Stage 1 Build A Minimal Agent Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f85d6ca-e5e9-4655-924d-afb9e5e0aafa",
   "metadata": {},
   "source": [
    "- [x] Use an LLM API to complete a basic conversation.\n",
    "- [x] Get the model to output structured JSON.\n",
    "- [x] Define a tool function, e.g. `search`, `calculator`, `read_file`.\n",
    "- [x] Parse the model's tool call / function call.\n",
    "- [x] Execute the tool and feed the result back to the model.\n",
    "- [x] Add max steps, timeout, and error handling to the agent loop.\n",
    "\n",
    "Docs:\n",
    "- [Function calling with the Gemini API](https://ai.google.dev/gemini-api/docs/function-calling?example=meeting)\n",
    "- [Google Gen AI SDK](https://googleapis.github.io/python-genai/#function-calling)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67560c5a-648b-445b-88c7-4a5fd46914fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 0. Setup\n",
    "from google import genai\n",
    "\n",
    "client = genai.Client(\n",
    "    vertexai=True,\n",
    "    project=\"xxxx\", # use yourself one\n",
    "    location=\"global\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e90e3122-5697-4f63-bb81-bca3786cc760",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It's currently **12:20 PM** on Friday, May 17, 2024, in Germany.\n",
      "\n",
      "Germany observes Central European Summer Time (CEST), which is UTC+2.\n"
     ]
    }
   ],
   "source": [
    "# 1. LLM API call\n",
    "\n",
    "model = \"gemini-2.5-flash\"\n",
    "prompts = r\"\"\"\n",
    "What's the time in Germany??\n",
    "\"\"\"\n",
    "\n",
    "response = client.models.generate_content(\n",
    "    model=model,\n",
    "    contents=prompts,\n",
    ")\n",
    "print(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7fd0e035-10ad-44d1-86e7-4167278acada",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time=datetime.datetime(2024, 5, 24, 12, 30, tzinfo=TzInfo(7200))\n"
     ]
    }
   ],
   "source": [
    "# 2. Structured outputs\n",
    "\n",
    "from pydantic import BaseModel\n",
    "from google.genai import types\n",
    "from datetime import datetime\n",
    "\n",
    "class TimeInfo(BaseModel):\n",
    "    time: datetime\n",
    "\n",
    "response = client.models.generate_content(\n",
    "    model=model,\n",
    "    contents=prompts,\n",
    "    config=types.GenerateContentConfig(\n",
    "        response_mime_type='application/json',\n",
    "        response_schema=TimeInfo,\n",
    "    ),\n",
    ")\n",
    "\n",
    "info = TimeInfo.model_validate_json(response.text)\n",
    "print(info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4ec24910-b808-420d-9836-54ba155f03cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Define a tool\n",
    "\n",
    "get_time_declaration = types.FunctionDeclaration(\n",
    "    name='get_time',\n",
    "    description='Give the current time given the time zone.',\n",
    "    parameters_json_schema={\n",
    "        'type': 'object',\n",
    "        'properties': {\n",
    "            'timezone': {\n",
    "                'type': 'string',\n",
    "                'description': 'time zone in pytz format, e.g. America/Los_Angeles',\n",
    "            }\n",
    "        },\n",
    "        'required': ['timezone'],\n",
    "    },\n",
    ")\n",
    "\n",
    "\n",
    "def get_time(timezone: str)-> datetime:\n",
    "    \"\"\"Get the current date and time of given the time zone.\n",
    "\n",
    "    Args:\n",
    "    timezone\n",
    "\n",
    "    Returns: datetime\n",
    "    \"\"\"\n",
    "    import pytz\n",
    "\n",
    "    timezone = pytz.timezone(timezone)\n",
    "    now = datetime.now(tz=timezone)\n",
    "    return now.strftime(\"%A, %d %B %Y, %H:%M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8e34441e-eeb3-409e-9d09-954263ac7db7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id=None args={'timezone': 'Europe/Berlin'} name='get_time' partial_args=None will_continue=None\n"
     ]
    }
   ],
   "source": [
    "# 4. Parse model's tool call\n",
    "\n",
    "# Configure the client and tools\n",
    "get_time_tool = types.Tool(function_declarations=[get_time_declaration])\n",
    "config = types.GenerateContentConfig(tools=[get_time_tool])\n",
    "\n",
    "# Define user prompt\n",
    "contents = [\n",
    "    types.Content(\n",
    "        role=\"user\", parts=[types.Part(text=prompts)]\n",
    "    )\n",
    "]\n",
    "\n",
    "# Send request with function declarations\n",
    "response = client.models.generate_content(\n",
    "    model=model,\n",
    "    contents=contents,\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "print(response.candidates[0].content.parts[0].function_call)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1265ac0c-8a61-4817-9cc4-837878416feb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Function execution result: Thursday, 28 May 2026, 18:12\n",
      "Clock Agent output:The current time in Germany is 18:12 on Thursday, 28 May 2026..\n"
     ]
    }
   ],
   "source": [
    "# 5. Excute the tool and feed back to model\n",
    "\n",
    "# Extract tool call details\n",
    "tool_call = response.candidates[0].content.parts[0].function_call\n",
    "\n",
    "\n",
    "if tool_call.name == \"get_time\":\n",
    "    result = get_time(**tool_call.args)\n",
    "    print(f\"Function execution result: {result}\")\n",
    "\n",
    "# Create a function response part\n",
    "function_response_part = types.Part.from_function_response(\n",
    "    name=tool_call.name,\n",
    "    response={\"result\": result},\n",
    "    # id=tool_call.id, used for Gemini 3.X\n",
    ")\n",
    "\n",
    "# Append function call and result of the function execution to contents\n",
    "contents.append(response.candidates[0].content) # Append the content from the model's response.\n",
    "contents.append(types.Content(role=\"user\", parts=[function_response_part])) # Append the function response\n",
    "\n",
    "final_response = client.models.generate_content(\n",
    "                    model=model,\n",
    "                    config=config,\n",
    "                    contents=contents,\n",
    "                    )\n",
    "\n",
    "print(f\"Clock Agent output:{final_response.text}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8d06a63c-dfbb-4fbe-8370-774f5c0906f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Clock Agent output: Could you please provide a more specific city in China?\n"
     ]
    }
   ],
   "source": [
    "# 6. Add max steps, timeout, and error handling to the agent loop\n",
    "\n",
    "# max step is to limit how many times we allow the model to use tools, to prevent model's endless tool calling and expensive bills\n",
    "\n",
    "import time\n",
    "\n",
    "TOOLS = {\"get_time\": get_time}\n",
    "prompts = r\"\"\"What's the time in China?\"\"\"\n",
    "contents = [\n",
    "    types.Content(\n",
    "        role=\"user\", parts=[types.Part(text=prompts)]\n",
    "    )\n",
    "]\n",
    "\n",
    "step = 0\n",
    "max_steps = 5\n",
    "timeout = 10  # seconds\n",
    "start_time = time.time()\n",
    "\n",
    "\n",
    "while step < max_steps:\n",
    "    # handle timeout\n",
    "    if time.time() - start_time > timeout:\n",
    "        print(\"Timeout reached, Stop!\")\n",
    "        break\n",
    "        \n",
    "    # Parse and excute the tool call\n",
    "    response = client.models.generate_content(\n",
    "        model=model,\n",
    "        contents=contents,\n",
    "        config=config,\n",
    "    )\n",
    "    tool_call_part = response.candidates[0].content.parts[0]\n",
    "    contents.append(response.candidates[0].content)\n",
    "    \n",
    "    if tool_call_part.function_call:\n",
    "        tc = tool_call_part.function_call\n",
    "        print(f\"[DEBUG] calling {tc.name} with args {tc.args}\")   # model return based on tool declarations\n",
    "        try:\n",
    "            result = TOOLS[tc.name](**tc.args)\n",
    "            print(f\"[DEBUG] tool result: {result}\")               # what does the tool return\n",
    "        except Exception as e:\n",
    "            print(f\"[DEBUG] tool FAILED: {type(e).__name__}: {e}\") # failed\n",
    "            result = f\"Error executing {tc.name}: {e}\"\n",
    "\n",
    "        \n",
    "        contents.append(types.Content(\n",
    "                role=\"user\",\n",
    "                parts=[types.Part.from_function_response(name=tc.name, response={\"result\": result})],\n",
    "            ))\n",
    "        step += 1 \n",
    "        continue\n",
    "    else:\n",
    "        print(f\"Clock Agent output: {response.text}\")\n",
    "        break\n",
    "\n",
    "if step >= max_steps:\n",
    "    print(\"MAX steps reached.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0ff61e31-39d7-4346-9cea-e67294e7020b",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}