microsoft--graphrag/packages/graphrag-llm/notebooks/10_tool_calling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a56845b0",
   "metadata": {},
   "source": [
    "# Function Tool Calling\n",
    "\n",
    "In order to use function tools, the completion endpoint needs a json schema of the function(s). This notebook uses `pydantic` to describe a function and its parameters and the `OpenAI` built-in `pydantic_function_tool` to create the necessary json schema. Other techniques may be used to create a definition for your functions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daf62482",
   "metadata": {},
   "source": [
    "## Manual Function Tool Calling\n",
    "\n",
    "This example demonstrates function tool calling by manually using `pydantic` and `pydantic_function_tool`. See the next example for a simplified approach.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53437ac4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding 5 and 7 gives you 12.\n"
     ]
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License\n",
    "\n",
    "import json\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from graphrag_llm.completion import LLMCompletion, create_completion\n",
    "from graphrag_llm.config import AuthMethod, ModelConfig\n",
    "from graphrag_llm.types import LLMCompletionResponse\n",
    "from graphrag_llm.utils import (\n",
    "    CompletionMessagesBuilder,\n",
    ")\n",
    "from openai import pydantic_function_tool\n",
    "from pydantic import BaseModel, ConfigDict, Field\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "api_key = os.getenv(\"GRAPHRAG_API_KEY\")\n",
    "model_config = ModelConfig(\n",
    "    model_provider=\"azure\",\n",
    "    model=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    azure_deployment_name=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    api_base=os.getenv(\"GRAPHRAG_API_BASE\"),\n",
    "    api_version=os.getenv(\"GRAPHRAG_API_VERSION\", \"2025-04-01-preview\"),\n",
    "    api_key=api_key,\n",
    "    auth_method=AuthMethod.AzureManagedIdentity if not api_key else AuthMethod.ApiKey,\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config)\n",
    "\n",
    "\n",
    "class AddTwoNumbers(BaseModel):\n",
    "    \"\"\"Input Argument for add two numbers.\"\"\"\n",
    "\n",
    "    model_config = ConfigDict(\n",
    "        extra=\"forbid\",\n",
    "    )\n",
    "\n",
    "    a: int = Field(description=\"The first number to add.\")\n",
    "    b: int = Field(description=\"The second number to add.\")\n",
    "\n",
    "\n",
    "# The actual function\n",
    "def add_two_numbers(options: AddTwoNumbers) -> int:\n",
    "    \"\"\"Add two numbers.\"\"\"\n",
    "    return options.a + options.b\n",
    "\n",
    "\n",
    "add_definition = pydantic_function_tool(\n",
    "    AddTwoNumbers,\n",
    "    # Function name and description\n",
    "    name=\"my_add_two_numbers_function\",\n",
    "    description=\"Add two numbers.\",\n",
    ")\n",
    "\n",
    "# Mapping of available functions\n",
    "available_functions = {\n",
    "    \"my_add_two_numbers_function\": {\n",
    "        \"function\": add_two_numbers,\n",
    "        \"input_model\": AddTwoNumbers,\n",
    "    },\n",
    "}\n",
    "\n",
    "messages_builder = CompletionMessagesBuilder().add_user_message(\n",
    "    \"Add 5 and 7 using a function call.\"\n",
    ")\n",
    "\n",
    "response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=messages_builder.build(),\n",
    "    tools=[add_definition],\n",
    ")  # type: ignore\n",
    "\n",
    "if not response.choices[0].message.tool_calls:\n",
    "    msg = \"No function call found in response.\"\n",
    "    raise ValueError(msg)\n",
    "\n",
    "# Add the assistant message with the function call to the message history\n",
    "messages_builder.add_assistant_message(\n",
    "    message=response.choices[0].message,\n",
    ")\n",
    "\n",
    "for tool_call in response.choices[0].message.tool_calls:\n",
    "    tool_id = tool_call.id\n",
    "    if tool_call.type != \"function\":\n",
    "        continue\n",
    "    function_name = tool_call.function.name\n",
    "    function_args = tool_call.function.arguments\n",
    "\n",
    "    args_dict = json.loads(function_args)\n",
    "\n",
    "    InputModel = available_functions[function_name][\"input_model\"]\n",
    "    function = available_functions[function_name][\"function\"]\n",
    "    input_options = InputModel(**args_dict)\n",
    "\n",
    "    result = function(input_options)\n",
    "\n",
    "    messages_builder.add_tool_message(\n",
    "        content=str(result),\n",
    "        tool_call_id=tool_id,\n",
    "    )\n",
    "\n",
    "final_response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=messages_builder.build(),\n",
    ")  # type: ignore\n",
    "print(final_response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31c7a9c",
   "metadata": {},
   "source": [
    "### Function Tool Definition\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eb6950e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"type\": \"function\",\n",
      "  \"function\": {\n",
      "    \"name\": \"my_add_two_numbers_function\",\n",
      "    \"strict\": true,\n",
      "    \"parameters\": {\n",
      "      \"additionalProperties\": false,\n",
      "      \"description\": \"Input Argument for add two numbers.\",\n",
      "      \"properties\": {\n",
      "        \"a\": {\n",
      "          \"description\": \"The first number to add.\",\n",
      "          \"title\": \"A\",\n",
      "          \"type\": \"integer\"\n",
      "        },\n",
      "        \"b\": {\n",
      "          \"description\": \"The second number to add.\",\n",
      "          \"title\": \"B\",\n",
      "          \"type\": \"integer\"\n",
      "        }\n",
      "      },\n",
      "      \"required\": [\n",
      "        \"a\",\n",
      "        \"b\"\n",
      "      ],\n",
      "      \"title\": \"AddTwoNumbers\",\n",
      "      \"type\": \"object\"\n",
      "    },\n",
      "    \"description\": \"Add two numbers.\"\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# View the output schema\n",
    "# This is what is passed to the completion tools param\n",
    "# Created using pydantic and pydantic_function_tool\n",
    "# but may be created manually as well\n",
    "print(json.dumps(add_definition, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "660de4c9",
   "metadata": {},
   "source": [
    "## Tool Calling with FunctionToolManager\n",
    "\n",
    "If using `pydantic` to describe function arguments, you can use the `FunctionToolManager` to register functions, produce defintions, and call functions in response to the LLM. This helps automate some of the above work.\n",
    "\n",
    "The following example demonstrates calling multiple functions in one LLM call.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fae701e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding numbers: 3 8\n",
      "Multiplying numbers: 9 5\n",
      "Reversing text: GraphRAG\n",
      "3 + 8 is 11, 9 * 5 is 45, and the reversed string 'GraphRAG' is 'GARhparG'.\n"
     ]
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License\n",
    "\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from graphrag_llm.completion import LLMCompletion, create_completion\n",
    "from graphrag_llm.config import AuthMethod, ModelConfig\n",
    "from graphrag_llm.types import LLMCompletionResponse\n",
    "from graphrag_llm.utils import (\n",
    "    CompletionMessagesBuilder,\n",
    "    FunctionToolManager,\n",
    ")\n",
    "from pydantic import BaseModel, ConfigDict, Field\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "api_key = os.getenv(\"GRAPHRAG_API_KEY\")\n",
    "model_config = ModelConfig(\n",
    "    model_provider=\"azure\",\n",
    "    model=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    azure_deployment_name=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    api_base=os.getenv(\"GRAPHRAG_API_BASE\"),\n",
    "    api_version=os.getenv(\"GRAPHRAG_API_VERSION\", \"2025-04-01-preview\"),\n",
    "    api_key=api_key,\n",
    "    auth_method=AuthMethod.AzureManagedIdentity if not api_key else AuthMethod.ApiKey,\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config)\n",
    "\n",
    "\n",
    "class NumbersInput(BaseModel):\n",
    "    \"\"\"Numbers input.\"\"\"\n",
    "\n",
    "    model_config = ConfigDict(\n",
    "        extra=\"forbid\",\n",
    "    )\n",
    "\n",
    "    a: int = Field(description=\"The first number.\")\n",
    "    b: int = Field(description=\"The second number.\")\n",
    "\n",
    "\n",
    "def add(options: NumbersInput) -> str:\n",
    "    \"\"\"Add two numbers.\"\"\"\n",
    "    # Print something to ensure function is called for verification\n",
    "    print(\"Adding numbers:\", options.a, options.b)\n",
    "    return str(options.a + options.b)\n",
    "\n",
    "\n",
    "def multiply(options: NumbersInput) -> str:\n",
    "    \"\"\"Multiply two numbers.\"\"\"\n",
    "    # Print something to ensure function is called for verification\n",
    "    print(\"Multiplying numbers:\", options.a, options.b)\n",
    "    return str(options.a * options.b)\n",
    "\n",
    "\n",
    "class TextInput(BaseModel):\n",
    "    \"\"\"Text input.\"\"\"\n",
    "\n",
    "    model_config = ConfigDict(\n",
    "        extra=\"forbid\",\n",
    "    )\n",
    "\n",
    "    test: str = Field(description=\"The string to reverse.\")\n",
    "\n",
    "\n",
    "def reverse_text(options: TextInput) -> str:\n",
    "    \"\"\"Reverse a string.\"\"\"\n",
    "    # Print something to ensure function is called for verification\n",
    "    print(\"Reversing text:\", options.test)\n",
    "    return options.test[::-1]\n",
    "\n",
    "\n",
    "function_tool_manager = FunctionToolManager()\n",
    "\n",
    "function_tool_manager.register_function_tool(\n",
    "    name=\"add\",\n",
    "    description=\"Add two numbers.\",\n",
    "    function=add,\n",
    "    input_model=NumbersInput,\n",
    ")\n",
    "function_tool_manager.register_function_tool(\n",
    "    name=\"multiply\",\n",
    "    description=\"Multiply two numbers.\",\n",
    "    function=multiply,\n",
    "    input_model=NumbersInput,\n",
    ")\n",
    "function_tool_manager.register_function_tool(\n",
    "    name=\"reverse_text\",\n",
    "    description=\"Reverse a string.\",\n",
    "    function=reverse_text,\n",
    "    input_model=TextInput,\n",
    ")\n",
    "\n",
    "\n",
    "messages_builder = CompletionMessagesBuilder().add_user_message(\n",
    "    \"What is 3 + 8 and 9 * 5? Also, reverse the string 'GraphRAG'.\"\n",
    ")\n",
    "\n",
    "# Multiple tool calls in parallel\n",
    "response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=messages_builder.build(),\n",
    "    tools=function_tool_manager.definitions(),\n",
    "    parallel_tool_calls=True,\n",
    ")  # type: ignore\n",
    "\n",
    "# Add the assistant message with the function call to the message history\n",
    "messages_builder.add_assistant_message(\n",
    "    message=response.choices[0].message,\n",
    ")\n",
    "\n",
    "tool_results = function_tool_manager.call_functions(response)\n",
    "\n",
    "for tool_message in tool_results:\n",
    "    messages_builder.add_tool_message(**tool_message)\n",
    "\n",
    "final_response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=messages_builder.build(),\n",
    ")  # type: ignore\n",
    "print(final_response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2d36f7a",
   "metadata": {},
   "source": [
    "## MCP Tools\n",
    "\n",
    "**Not currently supported**. `graphrag_llm` currently only implements the `completion` endpoints which do not support MCP tools.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}