{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ce4a9ccf-4bd6-43fb-a24d-b6a7da401a96",
   "metadata": {},
   "source": [
    "## Load xLAM model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1351d81-4502-4b65-b88a-464acd0e80f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "torch.random.manual_seed(0) \n",
    "\n",
    "model_name = \"Salesforce/xLAM-7b-r\"\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\", torch_dtype=\"auto\", trust_remote_code=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cdd5bae-da43-4713-9956-360f1f3a9721",
   "metadata": {},
   "source": [
    "## Build the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e138e9f6-0543-427c-bce6-b4f14765a040",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Please use our provided instruction prompt for best performance\n",
    "task_instruction = \"\"\"\n",
    "Based on the previous context and API request history, generate an API request or a response as an AI assistant.\"\"\".strip()\n",
    "\n",
    "format_instruction = \"\"\"\n",
    "The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
    "tool_calls an empty list \"[]\".\n",
    "```\n",
    "{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
    "```\n",
    "\"\"\".strip()\n",
    "\n",
    "get_weather_api = {\n",
    "    \"name\": \"get_weather\",\n",
    "    \"description\": \"Get the current weather for a location\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"location\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city and state, e.g. San Francisco, New York\"\n",
    "            },\n",
    "            \"unit\": {\n",
    "                \"type\": \"string\",\n",
    "                \"enum\": [\"celsius\", \"fahrenheit\"],\n",
    "                \"description\": \"The unit of temperature to return\"\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"location\"]\n",
    "    }\n",
    "}\n",
    "\n",
    "search_api = {\n",
    "    \"name\": \"search\",\n",
    "    \"description\": \"Search for information on the internet\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"query\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The search query, e.g. 'latest news on AI'\"\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"query\"]\n",
    "    }\n",
    "}\n",
    "\n",
    "openai_format_tools = [get_weather_api, search_api]\n",
    "\n",
    "# Define the input query and available tools\n",
    "query = \"What's the weather like in New York in fahrenheit?\"\n",
    "\n",
    "# Helper function to convert openai format tools to our more concise xLAM format\n",
    "def convert_to_xlam_tool(tools):\n",
    "    ''''''\n",
    "    if isinstance(tools, dict):\n",
    "        return {\n",
    "            \"name\": tools[\"name\"],\n",
    "            \"description\": tools[\"description\"],\n",
    "            \"parameters\": {k: v for k, v in tools[\"parameters\"].get(\"properties\", {}).items()}\n",
    "        }\n",
    "    elif isinstance(tools, list):\n",
    "        return [convert_to_xlam_tool(tool) for tool in tools]\n",
    "    else:\n",
    "        return tools\n",
    "\n",
    "def build_conversation_history_prompt(conversation_history: str):\n",
    "    parsed_history = []\n",
    "    for step_data in conversation_history:\n",
    "        parsed_history.append({\n",
    "            \"step_id\": step_data[\"step_id\"],\n",
    "            \"thought\": step_data[\"thought\"],\n",
    "            \"tool_calls\": step_data[\"tool_calls\"],\n",
    "            \"next_observation\": step_data[\"next_observation\"],\n",
    "            \"user_input\": step_data['user_input']\n",
    "        })\n",
    "        \n",
    "    history_string = json.dumps(parsed_history)\n",
    "    return f\"\\n[BEGIN OF HISTORY STEPS]\\n{history_string}\\n[END OF HISTORY STEPS]\\n\"\n",
    "    \n",
    "    \n",
    "# Helper function to build the input prompt for our model\n",
    "def build_prompt(task_instruction: str, format_instruction: str, tools: list, query: str, conversation_history: list):\n",
    "    prompt = f\"[BEGIN OF TASK INSTRUCTION]\\n{task_instruction}\\n[END OF TASK INSTRUCTION]\\n\\n\"\n",
    "    prompt += f\"[BEGIN OF AVAILABLE TOOLS]\\n{json.dumps(xlam_format_tools)}\\n[END OF AVAILABLE TOOLS]\\n\\n\"\n",
    "    prompt += f\"[BEGIN OF FORMAT INSTRUCTION]\\n{format_instruction}\\n[END OF FORMAT INSTRUCTION]\\n\\n\"\n",
    "    prompt += f\"[BEGIN OF QUERY]\\n{query}\\n[END OF QUERY]\\n\\n\"\n",
    "    \n",
    "    if len(conversation_history) > 0: prompt += build_conversation_history_prompt(conversation_history)\n",
    "    return prompt\n",
    "\n",
    "\n",
    "    \n",
    "# Build the input and start the inference\n",
    "xlam_format_tools = convert_to_xlam_tool(openai_format_tools)\n",
    "\n",
    "conversation_history = []\n",
    "content = build_prompt(task_instruction, format_instruction, xlam_format_tools, query, conversation_history)\n",
    "\n",
    "messages=[\n",
    "    { 'role': 'user', 'content': content}\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ff7bccd5-fa04-4fbe-92b3-13f58914da4d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[BEGIN OF TASK INSTRUCTION]\n",
      "Based on the previous context and API request history, generate an API request or a response as an AI assistant.\n",
      "[END OF TASK INSTRUCTION]\n",
      "\n",
      "[BEGIN OF AVAILABLE TOOLS]\n",
      "[{\"name\": \"get_weather\", \"description\": \"Get the current weather for a location\", \"parameters\": {\"location\": {\"type\": \"string\", \"description\": \"The city and state, e.g. San Francisco, New York\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit of temperature to return\"}}}, {\"name\": \"search\", \"description\": \"Search for information on the internet\", \"parameters\": {\"query\": {\"type\": \"string\", \"description\": \"The search query, e.g. 'latest news on AI'\"}}}]\n",
      "[END OF AVAILABLE TOOLS]\n",
      "\n",
      "[BEGIN OF FORMAT INSTRUCTION]\n",
      "The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
      "tool_calls an empty list \"[]\".\n",
      "```\n",
      "{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
      "```\n",
      "[END OF FORMAT INSTRUCTION]\n",
      "\n",
      "[BEGIN OF QUERY]\n",
      "What's the weather like in New York in fahrenheit?\n",
      "[END OF QUERY]\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5fb0006-9f5d-4d79-a8cd-819bad627441",
   "metadata": {},
   "source": [
    "## Get the model output (agent_action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbe56588-c786-4913-9062-373a22a92e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "# tokenizer.eos_token_id is the id of <|EOT|> token\n",
    "outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)\n",
    "agent_action = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b20ed2ae-86f6-489b-ad54-fe7ea911667b",
   "metadata": {},
   "source": [
    "For demo purpose, we use an example agent_action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ab20c084-44fa-403d-92a5-1b8ced72e9be",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent_action = \"\"\"{\"thought\": \"\", \"tool_calls\": [{\"name\": \"get_weather\", \"arguments\": {\"location\": \"New York\"}}]}\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd4d8e4-ee6b-499e-b75f-a48df7848a60",
   "metadata": {},
   "source": [
    "### Add follow-up question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "825649ba-2691-43a2-b3d8-7baf8b66d46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_agent_action(agent_action: str):\n",
    "    \"\"\"\n",
    "    Given an agent's action, parse it to add to conversation history\n",
    "    \"\"\"\n",
    "    try: parsed_agent_action_json = json.loads(agent_action)\n",
    "    except: return \"\", []\n",
    "    \n",
    "    if \"thought\" not in parsed_agent_action_json.keys(): thought = \"\"\n",
    "    else: thought = parsed_agent_action_json[\"thought\"]\n",
    "    \n",
    "    if \"tool_calls\" not in parsed_agent_action_json.keys(): tool_calls = []\n",
    "    else: tool_calls = parsed_agent_action_json[\"tool_calls\"]\n",
    "    \n",
    "    return thought, tool_calls\n",
    "\n",
    "def update_conversation_history(conversation_history: list, agent_action: str, environment_response: str, user_input: str):\n",
    "    \"\"\"\n",
    "    Update the conversation history list based on the new agent_action, environment_response, and/or user_input\n",
    "    \"\"\"\n",
    "    thought, tool_calls = parse_agent_action(agent_action)\n",
    "    new_step_data = {\n",
    "        \"step_id\": len(conversation_history) + 1,\n",
    "        \"thought\": thought,\n",
    "        \"tool_calls\": tool_calls,\n",
    "        \"next_observation\": environment_response,\n",
    "        \"user_input\": user_input,\n",
    "    }\n",
    "    \n",
    "    conversation_history.append(new_step_data)\n",
    "\n",
    "def get_environment_response(agent_action: str):\n",
    "    \"\"\"\n",
    "    Get the environment response for the agent_action\n",
    "    \"\"\"\n",
    "    # TODO: add custom implementation here\n",
    "    error_message, response_message = \"\", \"Sunny, 81 degrees\"\n",
    "    return {\"error\": error_message, \"response\": response_message}\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "051e6aff-c21b-4dcb-9eb8-c34154d90c39",
   "metadata": {},
   "source": [
    "1. **Get the next state after agent's response:**\n",
    "  The next 2 lines are examples of getting environment response and user_input.\n",
    "  It is depended on particular usage, we can have either one or both of those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "649a8e9d-9757-408c-9214-0590556c2db4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "environment_response = get_environment_response(agent_action)\n",
    "user_input = \"Now, search on the Internet for cute puppies\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c9c9418-1c54-4381-81d1-7f3834037739",
   "metadata": {},
   "source": [
    "2. After we got environment_response and (or) user_input, we want to add to our conversation history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bcfe89f3-8237-41bf-b92c-7c7568366042",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'step_id': 1,\n",
       "  'thought': '',\n",
       "  'tool_calls': [{'name': 'get_weather',\n",
       "    'arguments': {'location': 'New York'}}],\n",
       "  'next_observation': {'error': '', 'response': 'Sunny, 81 degrees'},\n",
       "  'user_input': 'Now, search on the Internet for cute puppies'}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "update_conversation_history(conversation_history, agent_action, environment_response, user_input)\n",
    "conversation_history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23ba97c6-2356-49e8-a07b-0e664b7f505c",
   "metadata": {},
   "source": [
    "3. We now can build the prompt with the updated history, and prepare the inputs for the LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ed204b3a-3be5-431b-b355-facaf31309d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "content = build_prompt(task_instruction, format_instruction, xlam_format_tools, query, conversation_history)\n",
    "messages=[\n",
    "    { 'role': 'user', 'content': content}\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8af843aa-6a47-4938-a455-567ea0cccce3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[BEGIN OF TASK INSTRUCTION]\n",
      "Based on the previous context and API request history, generate an API request or a response as an AI assistant.\n",
      "[END OF TASK INSTRUCTION]\n",
      "\n",
      "[BEGIN OF AVAILABLE TOOLS]\n",
      "[{\"name\": \"get_weather\", \"description\": \"Get the current weather for a location\", \"parameters\": {\"location\": {\"type\": \"string\", \"description\": \"The city and state, e.g. San Francisco, New York\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit of temperature to return\"}}}, {\"name\": \"search\", \"description\": \"Search for information on the internet\", \"parameters\": {\"query\": {\"type\": \"string\", \"description\": \"The search query, e.g. 'latest news on AI'\"}}}]\n",
      "[END OF AVAILABLE TOOLS]\n",
      "\n",
      "[BEGIN OF FORMAT INSTRUCTION]\n",
      "The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
      "tool_calls an empty list \"[]\".\n",
      "```\n",
      "{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
      "```\n",
      "[END OF FORMAT INSTRUCTION]\n",
      "\n",
      "[BEGIN OF QUERY]\n",
      "What's the weather like in New York in fahrenheit?\n",
      "[END OF QUERY]\n",
      "\n",
      "\n",
      "[BEGIN OF HISTORY STEPS]\n",
      "[{\"step_id\": 1, \"thought\": \"\", \"tool_calls\": [{\"name\": \"get_weather\", \"arguments\": {\"location\": \"New York\"}}], \"next_observation\": {\"error\": \"\", \"response\": \"Sunny, 81 degrees\"}, \"user_input\": \"Now, search on the Internet for cute puppies\"}]\n",
      "[END OF HISTORY STEPS]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71f76a10-a152-49d7-aa6f-3060cc49b935",
   "metadata": {},
   "source": [
    "## Get the model output for follow-up question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30af06fd-4aa7-4550-af39-3a77b5951882",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
    "# 5. Generate the outputs & decode\n",
    "#   tokenizer.eos_token_id is the id of <|EOT|> token\n",
    "outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)\n",
    "agent_action = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}