example_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cdb3a793-de9d-4e31-8374-1d0963ece7b2",
   "metadata": {},
   "source": [
    "# Open LLM agents with schema-guided generation of function calls\n",
    "\n",
    "LLM agents use large language models (LLMs) to decide which tools to use for interacting with their environment during multi-step task solving. Tool usage is done via function calling. The LLM is prompted or fine-tuned to generate JSON output representing one or more function calls. To ensure that the JSON output of an LLM follows a tool-specific schema one can use constrained decoding methods to control next token generation.\n",
    "\n",
    "Here, [grammar-based sampling](https://github.com/ggerganov/llama.cpp/pull/1773) available in llama.cpp is used for constrained decoding. A JSON schema of available tools is converted to a grammar used for generating instances of that schema (= tool calls) during decoding. This is used for implementing a function calling interface for a Llama-2-70b model, an LLM with (limited) tool usage capabilities. \n",
    "\n",
    "The [implementation](https://github.com/krasserm/grammar-based-agents) uses LangChain interfaces and is compatible LangChain's [agent framework](https://python.langchain.com/docs/modules/agents/). In its current state, it is a simple prototype for demonstrating schema-guided generation in LangChain agents. It is general enough to be used with many other language models supported by llama.cpp, after some tweaks to the prompt templates.\n",
    "\n",
    "## Agent\n",
    "\n",
    "A Llama-2 agent with schema-guided generation of function calls can be created as follows. Further details are described in section [Components](#components)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "19c9218b-b6e2-4133-970e-ef1578840a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor  \n",
    "from langchain.tools import StructuredTool  \n",
    "from langchain_experimental.chat_models.llm_wrapper import Llama2Chat\n",
    "\n",
    "from gba.agent import Agent\n",
    "from gba.llm import LlamaCppClient  \n",
    "from gba.math import Llama2Math\n",
    "from gba.tool import ToolCalling\n",
    "\n",
    "from example_tools import (\n",
    "    create_event,  \n",
    "    search_images,  \n",
    "    search_internet,  \n",
    ")\n",
    "\n",
    "# Custom LangChain LLM implementation that interacts with a model \n",
    "# hosted on a llama.cpp server\n",
    "llm = LlamaCppClient(url=\"http://localhost:8080/completion\", temperature=-1)  \n",
    "  \n",
    "# Converts incoming messages into a Llama-2 compatible chat prompt \n",
    "# and implements the LangChain chat model interface\n",
    "chat_model = Llama2Chat(llm=llm)  \n",
    "  \n",
    "# Layers a tool calling protocol on top of chat_model resulting in\n",
    "# an interface similar to the ChatOpenAI function calling interface\n",
    "tool_calling_model = ToolCalling(model=chat_model)\n",
    "\n",
    "# Proxy for a math LLM hosted on a llmama.cpp server. Backend of the\n",
    "# `calculate` tool.\n",
    "math_llm = LlamaCppClient(url=\"http://localhost:8088/completion\", temperature=-1)\n",
    "math_model = Llama2Math(llm=math_llm)\n",
    "\n",
    "def calculate(expression: str):\n",
    "    \"\"\"A tool for evaluating mathematical expressions. Do not use variables in expression.\"\"\"\n",
    "    return math_model.run(message=expression)\n",
    "\n",
    "# List of tools created from Python functions and used by the agent  \n",
    "tools = [  \n",
    "    StructuredTool.from_function(search_internet),  \n",
    "    StructuredTool.from_function(search_images),  \n",
    "    StructuredTool.from_function(create_event),  \n",
    "    StructuredTool.from_function(calculate),  \n",
    "]\n",
    " \n",
    "# Custom LangChain agent implementation (similar to OpenAIFunctionsAgent)\n",
    "agent_obj = Agent.from_llm_and_tools(\n",
    "    model=tool_calling_model, \n",
    "    tools=tools,\n",
    ")\n",
    "\n",
    "# LangChain's AgentExecutor for running the agent loop\n",
    "agent = AgentExecutor.from_agent_and_tools(\n",
    "    agent=agent_obj, \n",
    "    tools=tools, \n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7e6c3d-b73f-4dd0-a7ec-add6067508c9",
   "metadata": {},
   "source": [
    "The `llm` is a proxy for a 4-bit quantized [Llama-2 70b chat model](https://huggingface.co/TheBloke/Llama-2-70B-Chat-GGUF) hosted on a [llama.cpp server](https://github.com/ggerganov/llama.cpp/blob/master/examples/server/README.md). Tools in this example are mainly mockups except the `calulate` tool which uses another LLM to interpret and evaluate mathematical queries. Instructions for serving these LLMs are available [here](https://github.com/krasserm/grammar-based-agents/blob/master/README.md#getting-started).\n",
    "\n",
    "### Tool usage\n",
    "\n",
    "Here's an example of a request that is decomposed by the agent into multiple steps, using a tool at each step. To increase the probability that the model generates an appropriate tool call at each step i.e. selects the right tool and arguments, an unconstrained thinking phase precedes the constrained tool call generation phase. This enables attention to thinking results during the tool call generation phase (details in section [ToolCalling](#ToolCalling))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c296707-dc0a-4d22-bd38-52c334e4b824",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\n",
      "Thoughts: I need to call the tool 'search_internet' to find out who Leonardo DiCaprio's current girlfriend is and then calculate her age raised to the 0.24 power using the tool 'calculate'.\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `search_internet` with `{'query': \"Leonardo DiCaprio's current girlfriend\"}`\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mLeonardo di Caprio started dating Vittoria Ceretti in 2023. She was born in Italy and is 25 years old\u001b[0m\n",
      "Thoughts: I need to call another tool to obtain more information, specifically the \"calculate\" tool, to calculate Vittoria Ceretti's age raised to the power of 0.24.\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `calculate` with `{'expression': '25^0.24'}`\n",
      "\u001b[0mExecuting Llama2Math Python code:\n",
      "```\n",
      "result = pow(25, 0.24)\n",
      "```\n",
      "\u001b[36;1m\u001b[1;3m2.16524\u001b[0m\n",
      "Thoughts: I have enough information to respond with a final answer to the user.\n",
      "\u001b[32;1m\u001b[1;3mVittoria Ceretti is 2.16524 years old.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Vittoria Ceretti is 2.16524 years old.'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.run(\"Who is Leonardo DiCaprio's current girlfriend and what is her age raised to the 0.24 power?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec5e1a09-d8a4-458f-acfc-17d3b62ea52a",
   "metadata": {},
   "source": [
    "In the last step, the agent responds to the user directly with the final answer. The answer `Vittoria Ceretti, 2.16524` is short but correct (at the time of writing).\n",
    "\n",
    "### Direct response\n",
    "\n",
    "The agent may also decide to respond to the user immediately, without prior tool usage. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a0d934ec-7192-40fc-b1e3-e72df0cde15a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\n",
      "Thoughts: I can directly respond to the user with a joke, here's one: \"Why don't scientists trust atoms? Because they make up everything!\"\n",
      "\u001b[32;1m\u001b[1;3mWhy don't scientists trust atoms? Because they make up everything!\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Why don't scientists trust atoms? Because they make up everything!\""
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.run(\"Tell me a joke\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dead5fea-1b79-4615-b6c7-cd59a8c4b837",
   "metadata": {},
   "source": [
    "### External memory\n",
    "\n",
    "For maintaining conversational state with the user, the agent can be configured with external `memory`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2422924c-4fd2-47a5-b70a-6c4313fb8e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import MessagesPlaceholder  \n",
    "\n",
    "chat_history = MessagesPlaceholder(variable_name=\"chat_history\")  \n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\", \n",
    "    return_messages=True,\n",
    ")\n",
    "\n",
    "# Agent that additionally maintains conversational state with the user\n",
    "conversational_agent_obj = Agent.from_llm_and_tools(\n",
    "    model=tool_calling_model, \n",
    "    tools=tools, \n",
    "    extra_prompt_messages=[chat_history],\n",
    ")  \n",
    "conversational_agent = AgentExecutor.from_agent_and_tools(\n",
    "    agent=conversational_agent_obj, \n",
    "    tools=tools, \n",
    "    memory=memory, \n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8f3ba25-e533-404c-a49b-a8e0b612ede8",
   "metadata": {},
   "source": [
    "Let's start the conversation by requesting an image of a brown dog."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3ab1cda4-28f5-4d54-8241-769802d3c5b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\n",
      "Thoughts: I need to call the tool search_images(query: str) to find an image of a brown dog.\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `search_images` with `{'query': 'brown dog'}`\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m[brown_dog_1.jpg](https://example.com/brown_dog_1.jpg)\u001b[0m\n",
      "Thoughts: I have enough information to respond with a final answer, which is the URL of the image file \"brown_dog_1.jpg\".\n",
      "\u001b[32;1m\u001b[1;3mHere is an image of a brown dog: <https://example.com/brown_dog_1.jpg>\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Here is an image of a brown dog: <https://example.com/brown_dog_1.jpg>'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversational_agent.run(\"find an image of a brown dog\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e4bc0bc-8d21-454a-9fe4-c882a3cf58ae",
   "metadata": {},
   "source": [
    "The next request is in context of the previous one, and the agent updates the search query accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "357a9730-d06f-4ccd-bf36-a4f2e0a6c437",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\n",
      "Thoughts: I need to call another tool to search for images of a running dog.\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `search_images` with `{'query': 'brown dog running'}`\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m[brown_dog_running_1.jpg](https://example.com/brown_dog_running_1.jpg)\u001b[0m\n",
      "Thoughts: I have enough information to respond with a final answer, which is the URL of the image of a brown dog running.\n",
      "\u001b[32;1m\u001b[1;3mhttps://example.com/brown_dog_running_1.jpg\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'https://example.com/brown_dog_running_1.jpg'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversational_agent.run(\"dog should be running too\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4b94fe-3bf9-4d64-b5f7-bb6ec81d5640",
   "metadata": {},
   "source": [
    "The agent was able to create the query `brown dog running` only because it had access to conversational state (`brown` is mentioned only in the first request).\n",
    "\n",
    "## Components"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d94f557e-d1ea-445c-8f0f-ad9af256420b",
   "metadata": {},
   "source": [
    "### `LlamaCppClient`\n",
    "\n",
    "`LlamaCppClient` is a proxy for a model hosted on a llama.cpp server. It implements LangChain's `LLM` interface and relies on the caller to provide a valid Llama-2 chat prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6a448e4f-f42e-45b0-9251-b61c1c93999c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Sure, here's a cute image of a brown dog for you! 🐶💩\\n\\n[Image description: A photo of a brown dog with a wagging tail, sitting on a green grassy field. The dog has a playful expression and is looking directly at the camera.]\\n\\nI hope this image brings a smile to your face! Is there anything else I can assist you with? 😊\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = \"<s>[INST] <<SYS>>\\nYou are a helpful assistant\\n<</SYS>>\\n\\nFind an image of brown dog [/INST]\"\n",
    "llm(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cddc07ba-97cb-458d-8d03-fba119da6f93",
   "metadata": {},
   "source": [
    "If a tool schema is provided, it is converted by `LlamaCppClient` to a grammar used for grammar-based sampling on the server side so that the output is a valid instance of that schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e9124bed-dcaa-44f3-96e9-cc01f6256408",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'object',\n",
       " 'properties': {'tool': {'const': 'search_images'},\n",
       "  'arguments': {'type': 'object',\n",
       "   'properties': {'query': {'type': 'string'}}}}}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from gba.tool import tool_to_schema\n",
    "\n",
    "tool_schema = tool_to_schema(StructuredTool.from_function(search_images))\n",
    "tool_schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "20be5559-0e09-4666-a0f1-dad033086ef4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{ \"arguments\": { \"query\": \"brown dog\" }, \"tool\": \"search_images\"}'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm(prompt, schema=tool_schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5ae585f-1137-4ee0-a8ef-85de4c5237d3",
   "metadata": {},
   "source": [
    "### `Llama2Chat`\n",
    "\n",
    "[Llama2Chat](https://python.langchain.com/docs/integrations/chat/llama2_chat) wraps `llm` and implements a chat model interface that applies the Llama-2 chat prompt to incoming messages (see also [these](https://github.com/langchain-ai/langchain/pull/8295#issuecomment-1668988543) [examples](https://github.com/langchain-ai/langchain/pull/8295#issuecomment-1811914445) for other chat prompt formats). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e360b6ed-a260-4a4b-a339-e3c9f808733c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Sure, here's a cute image of a brown dog for you! 🐶💩\\n\\n[Image description: A photo of a brown dog with a wagging tail, sitting on a green grassy field. The dog has a playful expression and is looking directly at the camera.]\\n\\nI hope this image brings a smile to your face! Is there anything else I can assist you with? 😊\")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.schema.messages import HumanMessage, SystemMessage\n",
    "\n",
    "messages = [\n",
    "    SystemMessage(content=\"You are a helpful assistant\"),\n",
    "    HumanMessage(content=\"Find an image of brown dog\"),\n",
    "]\n",
    "chat_model.predict_messages(messages=messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "61baa00e-6e67-40d3-a015-b9b83c130f90",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='{ \"arguments\": { \"query\": \"brown dog\" }, \"tool\": \"search_images\"}')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_model.predict_messages(messages=messages, schema=tool_to_schema(StructuredTool.from_function(search_images)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6214fe7f-381d-4f49-8b5a-6f4dd57a8952",
   "metadata": {},
   "source": [
    "### `ToolCalling`\n",
    "\n",
    "`ToolCalling` adds tool calling functionality to the wrapped `chat_model`, resulting in an interface very similar to the function calling interface of `ChatOpenAI`. It converts a list of provided tools to a [oneOf](https://json-schema.org/understanding-json-schema/reference/combining#oneof) JSON schema and submits it to the wrapped chat model. It also constructs a system prompt internally that informs the backend LLM about the presence of these tools and adds a special `respond_to_user` tool used by the LLM for providing a final response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "11bd32fd-3e89-4aa8-bf0c-d8e2804071c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Thoughts: I need to call the calculate tool to evaluate the expression and obtain the result.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AIMessage(content='', additional_kwargs={'tool_call': {'tool': 'calculate', 'arguments': {'expression': '(2 * 5) ** 0.8 / 2'}}})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat = [HumanMessage(content=\"What is (2 * 5) raised to power of 0.8 divided by 2?\")]\n",
    "tool_calling_model.predict_messages(messages=chat, tools=tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "759afbaf-6781-45ac-85c8-e4f4b2237078",
   "metadata": {},
   "source": [
    "The wrapped `chat_model` first receives a message with an instruction to think about the current state and what to do next. This message doesn't contain a schema so that the model output is unconstrained. Then the model receives another message with an instruction to act i.e. respond with a tool call. This message contains the JSON schema of the provided tools so that the tool call can be generated with grammar-based sampling."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b30041-e10d-4acf-bf9c-a83cfd84fa0c",
   "metadata": {},
   "source": [
    "## Outlook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c17faa46-9da6-4819-962f-bfe347fa270e",
   "metadata": {},
   "source": [
    "This two-phase approach of unconstrained thinking and constrained tool call generation may also be the basis for using more specialized models e.g. one for the thinking phase and another one for the tool call generation phase. I plan to present this in another article."
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}