a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
361 lines
9.3 KiB
Plaintext
361 lines
9.3 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "dcb12d1f",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Function Calling Google Gemini Agent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cb1eb6c7",
|
|
"metadata": {},
|
|
"source": [
|
|
"This notebook shows you how to use Google Gemini Agent, powered by function calling capabilities.\n",
|
|
"\n",
|
|
"Google's Gemini 2.5 Pro, Gemini 2.5 Flash, Gemini 2.5 Flash-Lite, Gemini 2.0 Flash models support function calling capabilities. You can find a comprehensive capabilities overview on the [model overview](https://ai.google.dev/gemini-api/docs/models) page."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f2f9042d",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initial Setup"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "9a377701",
|
|
"metadata": {},
|
|
"source": [
|
|
"Let's start by importing some simple building blocks.\n",
|
|
"\n",
|
|
"The main thing we need is:\n",
|
|
"\n",
|
|
"1. the Google Gemini API (using our own llama_index LLM class)\n",
|
|
"2. a place to keep conversation history\n",
|
|
"3. a definition for tools that our agent can use.\n",
|
|
"\n",
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "30218006",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-llms-google-genai llama-index -q"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d3a847cb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# import os\n",
|
|
"\n",
|
|
"# os.environ[\"GOOGLE_API_KEY\"] = \"...\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e61bf2cb",
|
|
"metadata": {},
|
|
"source": [
|
|
"Let's define some very simple calculator tools for our agent."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f7803c96",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def multiply(a: int, b: int) -> int:\n",
|
|
" \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
|
|
" return a * b\n",
|
|
"\n",
|
|
"\n",
|
|
"def add(a: int, b: int) -> int:\n",
|
|
" \"\"\"Add two integers and returns the result integer\"\"\"\n",
|
|
" return a + b"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "bd7bfda0",
|
|
"metadata": {},
|
|
"source": [
|
|
"Make sure your GOOGLE_API_KEY is set. Otherwise explicitly specify the api_key parameter."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "98e844a2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.llms.google_genai import GoogleGenAI\n",
|
|
"from google.genai import types\n",
|
|
"\n",
|
|
"llm = GoogleGenAI(\n",
|
|
" model=\"gemini-2.5-flash\",\n",
|
|
" generation_config=types.GenerateContentConfig(\n",
|
|
" thinking_config=types.ThinkingConfig(\n",
|
|
" thinking_budget=0\n",
|
|
" ) # Disables thinking\n",
|
|
" ),\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f5879b71",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initialize Google Gemini Agent\n",
|
|
"\n",
|
|
"Here we initialize a simple Google Gemini Agent agent with calculator functions."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fa419b7e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.agent.workflow import FunctionAgent\n",
|
|
"\n",
|
|
"agent = FunctionAgent(\n",
|
|
" tools=[multiply, add],\n",
|
|
" llm=llm,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3ba0e6c0",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.agent.workflow import ToolCallResult\n",
|
|
"\n",
|
|
"\n",
|
|
"async def run_agent_verbose(query: str):\n",
|
|
" handler = agent.run(query)\n",
|
|
" async for event in handler.stream_events():\n",
|
|
" if isinstance(event, ToolCallResult):\n",
|
|
" print(\n",
|
|
" f\"Called tool {event.tool_name} with args {event.tool_kwargs}\\nGot result: {event.tool_output}\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" return await handler"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0c430f5c",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Chat"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "95d89496",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Called tool add with args {'b': 2, 'a': 121}\n",
|
|
"Got result: 123\n",
|
|
"Called tool multiply with args {'a': 123, 'b': 5}\n",
|
|
"Got result: 615\n",
|
|
"The answer is 615.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"response = await run_agent_verbose(\"What is (121 + 2) * 5?\")\n",
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ce750c0e",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"[ToolCallResult(tool_name='add', tool_kwargs={'b': 2, 'a': 121}, tool_id='add', tool_output=ToolOutput(content='123', tool_name='add', raw_input={'args': (), 'kwargs': {'b': 2, 'a': 121}}, raw_output=123, is_error=False), return_direct=False), ToolCallResult(tool_name='multiply', tool_kwargs={'a': 123, 'b': 5}, tool_id='multiply', tool_output=ToolOutput(content='615', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 123, 'b': 5}}, raw_output=615, is_error=False), return_direct=False)]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# inspect sources\n",
|
|
"print(response.tool_calls)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ef7df617",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Managing Context/Memory\n",
|
|
"\n",
|
|
"By default, `.run()` is stateless. If you want to maintain state, you can pass in a context object."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ce6c27e2",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Your name is John Doe.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from llama_index.core.workflow import Context\n",
|
|
"\n",
|
|
"agent = FunctionAgent(llm=llm)\n",
|
|
"ctx = Context(agent)\n",
|
|
"\n",
|
|
"response = await agent.run(\"My name is John Doe\", ctx=ctx)\n",
|
|
"response = await agent.run(\"What is my name?\", ctx=ctx)\n",
|
|
"\n",
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "809b4158",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Google Gemini Agent over RAG Pipeline"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "5d99032f",
|
|
"metadata": {},
|
|
"source": [
|
|
"Build a Anthropic agent over a simple 10K document. We use OpenAI embeddings and Gemini 2.0 Flash to construct the RAG pipeline, and pass it to the Gemini 2.5 Flash agent as a tool."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "74c8908a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!mkdir -p 'data/10k/'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3bd5fbbb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.tools import QueryEngineTool\n",
|
|
"from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
|
|
"from llama_index.embeddings.openai import OpenAIEmbedding\n",
|
|
"from llama_index.llms.google_genai import GoogleGenAI\n",
|
|
"\n",
|
|
"embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-large\")\n",
|
|
"query_llm = GoogleGenAI(model=\"gemini-2.0-flash\")\n",
|
|
"\n",
|
|
"# load data\n",
|
|
"uber_docs = SimpleDirectoryReader(\n",
|
|
" input_files=[\"./data/10k/uber_2021.pdf\"]\n",
|
|
").load_data()\n",
|
|
"\n",
|
|
"# build index\n",
|
|
"uber_index = VectorStoreIndex.from_documents(\n",
|
|
" uber_docs, embed_model=embed_model\n",
|
|
")\n",
|
|
"uber_engine = uber_index.as_query_engine(similarity_top_k=3, llm=query_llm)\n",
|
|
"query_engine_tool = QueryEngineTool.from_defaults(\n",
|
|
" query_engine=uber_engine,\n",
|
|
" name=\"uber_10k\",\n",
|
|
" description=(\n",
|
|
" \"Provides information about Uber financials for year 2021. \"\n",
|
|
" \"Use a detailed plain text question as input to the tool.\"\n",
|
|
" ),\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ce39f025",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.agent.workflow import FunctionAgent\n",
|
|
"\n",
|
|
"agent = FunctionAgent(tools=[query_engine_tool], llm=llm, verbose=True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9f37f9f6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = await agent.run(\n",
|
|
" \"Tell me both the risk factors and tailwinds for Uber?\"\n",
|
|
")\n",
|
|
"print(str(response))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "gsoc",
|
|
"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"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|