wehub-resource-sync
2132 lines
90 KiB
Plaintext
2132 lines
90 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "f705f4be70e9"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright 2025 Google LLC\n",
|
|
"#\n",
|
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
|
"# you may not use this file except in compliance with the License.\n",
|
|
"# You may obtain a copy of the License at\n",
|
|
"#\n",
|
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
|
"#\n",
|
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
|
"# See the License for the specific language governing permissions and\n",
|
|
"# limitations under the License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2979962cdda2"
|
|
},
|
|
"source": [
|
|
"# Function Calling Agent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "c7d5e38ade6b"
|
|
},
|
|
"source": [
|
|
"<table align=\"left\">\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/agent-platform/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagents%2Fgenai-experience-concierge%2Fagent-design-patterns%2Ffunction-calling.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/agent-platform/workbench/instances?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/workbench-icon.svg\" alt=\"Workbench logo\"><br> Open in Workbench\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
"</table>\n",
|
|
"\n",
|
|
"<div style=\"clear: both;\"></div>\n",
|
|
"\n",
|
|
"<p>\n",
|
|
"<b>Share to:</b>\n",
|
|
"\n",
|
|
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agents/genai-experience-concierge/agent-design-patterns/function-calling.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
|
"</a>\n",
|
|
"</p>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "845489805d6f"
|
|
},
|
|
"source": [
|
|
"| | | |\n",
|
|
"|-|-|-|\n",
|
|
"|Author(s) | [Pablo Gaeta](https://github.com/pablofgaeta) | [Aadila Jasmin](https://github.com/aadi1405) |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "f5efbabd0ce2"
|
|
},
|
|
"source": [
|
|
"## Overview\n",
|
|
"\n",
|
|
"### Introduction\n",
|
|
"\n",
|
|
"Large Language Models (LLMs) are powerful at solving many types of problems. However, they are constrained by the following limitations:\n",
|
|
"\n",
|
|
"- They are frozen after training, leading to stale knowledge.\n",
|
|
"- They can't query or modify external data.\n",
|
|
"\n",
|
|
"Function calling can address these shortcomings. Function calling is sometimes referred to as tool use because it allows the model to use external tools such as APIs and functions.\n",
|
|
"\n",
|
|
"Function calling is a popular technique for enabling structured retrieval augmented generation (RAG) and enabling LLMs to take actions in the real world. Learn more about function calling [here](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling).\n",
|
|
"\n",
|
|
"This notebook constructs an agent with access to a collection of function declarations to search over a synthetic BigQuery dataset for a fictional company named \"Cymbal Retail\". The dataset contains information about products, store locations, and product-store inventory. The function declarations allow for structured query generation to enable the LLM to query the database in a secure, controlled manner. This approach can be contrasted with Natural-Language-To-SQL (NL2SQL) which can generate and execute arbitrary SQL, making it more flexible but more prone to security risks ([learn more about NL2SQL](https://cloud.google.com/blog/products/data-analytics/nl2sql-with-bigquery-and-gemini)). In addition to exact filtering mechanisms like setting a maximum product price or store search radius, the demo utilizes integrated BQML embedding support ([reference](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-embedding#text-embedding)) to re-rank results using product name/description semantic similarity.\n",
|
|
"\n",
|
|
"### Retail Search Assistant demo use cases\n",
|
|
"\n",
|
|
"1. Store Search. Filter by:\n",
|
|
" * Store Name\n",
|
|
" * Search Radius\n",
|
|
" * Product IDs\n",
|
|
" * Number of Results\n",
|
|
"\n",
|
|
"1. Product Search. Filter and rank by:\n",
|
|
" * Store IDs\n",
|
|
" * Price Range\n",
|
|
" * Number of Results\n",
|
|
" * Product Name/Description Semantic Similarity\n",
|
|
"\n",
|
|
"1. Inventory Check. Given a product and store, lookup current inventory.\n",
|
|
"\n",
|
|
"### Key Components\n",
|
|
"\n",
|
|
"The agent's architecture consists of the following components:\n",
|
|
"\n",
|
|
"* **Language Model:** Gemini is used for natural language understanding, function calling, and response generation.\n",
|
|
"* **State Management:** LangGraph manages the conversation flow and maintains the session state, including conversation history and user location.\n",
|
|
"* **Grounding Data:** BigQuery stores the Cymbal Retail data, including store information, product details, and inventory levels.\n",
|
|
"* **Tools:**\n",
|
|
" * `find_products`: Searches for products based on keywords, store filters, and price range.\n",
|
|
" * `find_stores`: Searches for stores based on location, name, and product availability.\n",
|
|
" * `find_inventory`: Retrieves the inventory level of a specific product at a given store.\n",
|
|
"\n",
|
|
"### Workflow\n",
|
|
"\n",
|
|
"1. Function calling handlers are generated for each defined tool based on runtime configurations (e.g. project, dataset, and user location).\n",
|
|
"1. Gemini is invoked with the function calling tools available.\n",
|
|
"1. If any functions are called, the calls and responses are streamed. Finally the text response is generated and streamed.\n",
|
|
"1. The conversation turn is finalized and the agent prepares for the next input."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "93b9c58f24d9"
|
|
},
|
|
"source": [
|
|
"## Get Started"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "d5027929de8f"
|
|
},
|
|
"source": [
|
|
"### Install dependencies"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "6911da456f3e"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install -q langgraph langgraph-checkpoint google-genai google-cloud-bigquery google-cloud-bigquery-storage db-dtypes thefuzz"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "f42d12d15616"
|
|
},
|
|
"source": [
|
|
"### Restart runtime\n",
|
|
"\n",
|
|
"To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
|
|
"\n",
|
|
"The restart might take a minute or longer. After it's restarted, continue to the next step."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "06fd78d27773"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# import IPython\n",
|
|
"\n",
|
|
"# app = IPython.Application.instance()\n",
|
|
"# app.kernel.do_shutdown(True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "e114f5653870"
|
|
},
|
|
"source": [
|
|
"### Authenticate your notebook environment (Colab only)\n",
|
|
"\n",
|
|
"If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {
|
|
"id": "911453311a5d"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import sys\n",
|
|
"\n",
|
|
"if \"google.colab\" in sys.modules:\n",
|
|
" from google.colab import auth\n",
|
|
"\n",
|
|
" auth.authenticate_user()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "d6388c6da021"
|
|
},
|
|
"source": [
|
|
"## Notebook parameters\n",
|
|
"\n",
|
|
"**Important:** In order to run this notebook, you must have a Cymbal Retail dataset created with a remote embedding model and product, stores, and inventory tables. These can be quickly provisioned by following the instructions at [Create the Cymbal Retail dataset](../README.md#optional-create-the-cymbal-retail-dataset)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {
|
|
"id": "daadb13ea48b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Use the environment variable if the user doesn't provide Project ID.\n",
|
|
"import os\n",
|
|
"\n",
|
|
"# fmt: off\n",
|
|
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
|
|
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
|
|
"\n",
|
|
"REGION = \"us-central1\" # @param {type:\"string\"}\n",
|
|
"CYMBAL_DATASET_LOCATION = \"US\" # @param {type:\"string\"}\n",
|
|
"CHAT_MODEL_NAME = \"gemini-3.5-flash\" # @param {type:\"string\"}\n",
|
|
"CYMBAL_DATASET = \"[project-id].[dataset-id]\" # @param {type:\"string\", placeholder: \"[project-id].[dataset-id]\"}\n",
|
|
"# fmt: on\n",
|
|
"\n",
|
|
"MAX_STORE_RESULTS = 10 # @param {type:\"integer\"}\n",
|
|
"MAX_PRODUCT_RESULTS = 10 # @param {type:\"integer\"}\n",
|
|
"STORE_NAME_SIMILARITY_THRESHOLD = 90 # @param {type:\"integer\"}\n",
|
|
"\n",
|
|
"CYMBAL_STORES_TABLE_URI = f\"{CYMBAL_DATASET}.cymbal_store\"\n",
|
|
"CYMBAL_PRODUCTS_TABLE_URI = f\"{CYMBAL_DATASET}.cymbal_product\"\n",
|
|
"CYMBAL_INVENTORY_TABLE_URI = f\"{CYMBAL_DATASET}.cymbal_inventory\"\n",
|
|
"CYMBAL_EMBEDDING_MODEL_URI = f\"{CYMBAL_DATASET}.text_embedding\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "413a89ab7cc5"
|
|
},
|
|
"source": [
|
|
"## Define the Function Calling Agent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "6a04ecba1630"
|
|
},
|
|
"source": [
|
|
"### Import dependencies"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {
|
|
"id": "12712546f708"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import asyncio\n",
|
|
"import datetime\n",
|
|
"import inspect\n",
|
|
"import json\n",
|
|
"import logging\n",
|
|
"import random\n",
|
|
"import uuid\n",
|
|
"from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping\n",
|
|
"from typing import Any, Literal, TypedDict\n",
|
|
"\n",
|
|
"import pydantic\n",
|
|
"from IPython import display as ipd\n",
|
|
"from google import genai\n",
|
|
"from google.cloud import bigquery\n",
|
|
"from google.genai import types as genai_types\n",
|
|
"from langchain_core.runnables import config as lc_config\n",
|
|
"from langgraph import graph\n",
|
|
"from langgraph import types as lg_types\n",
|
|
"from langgraph.checkpoint import memory as memory_checkpoint\n",
|
|
"from langgraph.config import get_stream_writer\n",
|
|
"from thefuzz import fuzz\n",
|
|
"\n",
|
|
"random.seed(0)\n",
|
|
"\n",
|
|
"logger = logging.getLogger(__name__)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "f6563a070bad"
|
|
},
|
|
"source": [
|
|
"### Define schemas\n",
|
|
"\n",
|
|
"Defines all of the schemas, constants, and types required for building the agent."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {
|
|
"id": "3a0f1257ed7b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"class AgentConfig(pydantic.BaseModel):\n",
|
|
" \"\"\"Configuration settings for the agent, including project, region, model, and data locations.\"\"\"\n",
|
|
"\n",
|
|
" project: str\n",
|
|
" \"\"\"The Google Cloud project ID.\"\"\"\n",
|
|
" region: str\n",
|
|
" \"\"\"The Google Cloud region where the agent is deployed.\"\"\"\n",
|
|
" chat_model_name: str\n",
|
|
" \"\"\"The name of the Gemini chat model to use.\"\"\"\n",
|
|
" cymbal_dataset_location: str\n",
|
|
" \"\"\"Location of the Cymbal dataset.\"\"\"\n",
|
|
" cymbal_products_table_uri: str\n",
|
|
" \"\"\"URI of the Cymbal products table.\"\"\"\n",
|
|
" cymbal_stores_table_uri: str\n",
|
|
" \"\"\"URI of the Cymbal stores table.\"\"\"\n",
|
|
" cymbal_inventory_table_uri: str\n",
|
|
" \"\"\"URI of the Cymbal inventory table.\"\"\"\n",
|
|
" cymbal_embedding_model_uri: str\n",
|
|
" \"\"\"URI of the Cymbal embedding model.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"# Node names and literal types\n",
|
|
"\n",
|
|
"CHAT_NODE_NAME = \"CHAT\"\n",
|
|
"\"\"\"The name of the chat node in the LangGraph.\"\"\"\n",
|
|
"\n",
|
|
"POST_PROCESS_NODE_NAME = \"POST_PROCESS\"\n",
|
|
"\"\"\"The name of the post-processing node in the LangGraph.\"\"\"\n",
|
|
"\n",
|
|
"PostProcessNodeTargetLiteral = Literal[\"POST_PROCESS\"]\n",
|
|
"\"\"\"Literal type for the post-processing node target.\"\"\"\n",
|
|
"\n",
|
|
"EndNodeTargetLiteral = Literal[\"__end__\"]\n",
|
|
"\"\"\"Literal type for the end node target.\"\"\"\n",
|
|
"\n",
|
|
"# DB Models\n",
|
|
"\n",
|
|
"\n",
|
|
"class Store(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents a store with its details.\"\"\"\n",
|
|
"\n",
|
|
" id: str\n",
|
|
" \"\"\"Unique identifier for the store.\"\"\"\n",
|
|
" name: str\n",
|
|
" \"\"\"Name of the store.\"\"\"\n",
|
|
" url: str\n",
|
|
" \"\"\"URL of the store's website.\"\"\"\n",
|
|
" street_address: str\n",
|
|
" \"\"\"Street address of the store.\"\"\"\n",
|
|
" city: str\n",
|
|
" \"\"\"City where the store is located.\"\"\"\n",
|
|
" state: str\n",
|
|
" \"\"\"State where the store is located.\"\"\"\n",
|
|
" zip_code: int\n",
|
|
" \"\"\"ZIP code of the store.\"\"\"\n",
|
|
" country: str\n",
|
|
" \"\"\"Country where the store is located.\"\"\"\n",
|
|
" phone_number: str\n",
|
|
" \"\"\"Phone number of the store.\"\"\"\n",
|
|
" latitude: float\n",
|
|
" \"\"\"Latitude of the store's location.\"\"\"\n",
|
|
" longitude: float\n",
|
|
" \"\"\"Longitude of the store's location.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"class Product(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents a product with its details.\"\"\"\n",
|
|
"\n",
|
|
" id: str\n",
|
|
" \"\"\"Unique identifier for the product.\"\"\"\n",
|
|
" name: str\n",
|
|
" \"\"\"Name of the product.\"\"\"\n",
|
|
" url: str\n",
|
|
" \"\"\"URL of the product's page.\"\"\"\n",
|
|
" description: str\n",
|
|
" \"\"\"Description of the product.\"\"\"\n",
|
|
" brand: str | None = None\n",
|
|
" \"\"\"Brand of the product (optional).\"\"\"\n",
|
|
" category: str\n",
|
|
" \"\"\"Category of the product.\"\"\"\n",
|
|
" available: bool\n",
|
|
" \"\"\"Availability status of the product.\"\"\"\n",
|
|
" list_price: float\n",
|
|
" \"\"\"List price of the product.\"\"\"\n",
|
|
" sale_price: float | None = None\n",
|
|
" \"\"\"Sale price of the product (optional).\"\"\"\n",
|
|
" currency: str = \"usd\"\n",
|
|
" \"\"\"Currency of the product's price.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"class Inventory(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents the inventory of a product in a store.\"\"\"\n",
|
|
"\n",
|
|
" store_id: str\n",
|
|
" \"\"\"Identifier of the store.\"\"\"\n",
|
|
" product_id: str\n",
|
|
" \"\"\"Identifier of the product.\"\"\"\n",
|
|
" value: int = 0\n",
|
|
" \"\"\"Quantity of the product in the store's inventory.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"# Tool return types\n",
|
|
"\n",
|
|
"\n",
|
|
"class StoreSearchResult(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents the result of a store search.\"\"\"\n",
|
|
"\n",
|
|
" stores: list[Store] = []\n",
|
|
" \"\"\"List of stores matching the search criteria.\"\"\"\n",
|
|
" query: str\n",
|
|
" \"\"\"The search query used.\"\"\"\n",
|
|
" error: str | None = None\n",
|
|
" \"\"\"Error message, if any.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"class ProductSearchResult(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents the result of a product search.\"\"\"\n",
|
|
"\n",
|
|
" products: list[Product] = []\n",
|
|
" \"\"\"List of products matching the search criteria.\"\"\"\n",
|
|
" query: str\n",
|
|
" \"\"\"The search query used.\"\"\"\n",
|
|
" error: str | None = None\n",
|
|
" \"\"\"Error message, if any.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"class InventorySearchResult(pydantic.BaseModel):\n",
|
|
" \"\"\"Represents the result of an inventory search.\"\"\"\n",
|
|
"\n",
|
|
" inventory: Inventory\n",
|
|
" \"\"\"The inventory information for the product in the store.\"\"\"\n",
|
|
" query: str\n",
|
|
" \"\"\"The search query used.\"\"\"\n",
|
|
" error: str | None = None\n",
|
|
" \"\"\"Error message, if any.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"# LangGraph models\n",
|
|
"\n",
|
|
"\n",
|
|
"class Turn(TypedDict, total=False):\n",
|
|
" \"\"\"Represents a single turn in a conversation.\n",
|
|
"\n",
|
|
" Attributes:\n",
|
|
" id: Unique identifier for the turn.\n",
|
|
" created_at: Timestamp of when the turn was created.\n",
|
|
" user_input: The user's input in this turn.\n",
|
|
" response: The agent's response in this turn, if any.\n",
|
|
" user_latitude: The user's latitude in this turn, if any.\n",
|
|
" user_longitude: The user's longitude in this turn, if any\n",
|
|
" messages: A list of Gemini content messages associated with this turn.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" id: uuid.UUID\n",
|
|
" \"\"\"Unique identifier for the turn.\"\"\"\n",
|
|
"\n",
|
|
" created_at: datetime.datetime\n",
|
|
" \"\"\"Timestamp of when the turn was created.\"\"\"\n",
|
|
"\n",
|
|
" user_input: str\n",
|
|
" \"\"\"The user's input for this turn.\"\"\"\n",
|
|
"\n",
|
|
" response: str\n",
|
|
" \"\"\"The agent's response for this turn, if any.\"\"\"\n",
|
|
"\n",
|
|
" user_latitude: float | None\n",
|
|
" \"\"\"The user's latitude for this turn, if any.\"\"\"\n",
|
|
"\n",
|
|
" user_longitude: float | None\n",
|
|
" \"\"\"The user's longitude for this turn, if any.\"\"\"\n",
|
|
"\n",
|
|
" messages: list[genai_types.Content]\n",
|
|
" \"\"\"List of Gemini Content objects representing the conversation messages in this turn.\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"class GraphSession(TypedDict, total=False):\n",
|
|
" \"\"\"Represents the complete state of a conversation session.\n",
|
|
"\n",
|
|
" Attributes:\n",
|
|
" id: Unique identifier for the session.\n",
|
|
" created_at: Timestamp of when the session was created.\n",
|
|
" current_turn: The current turn in the session, if any.\n",
|
|
" turns: A list of all turns in the session.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" id: uuid.UUID\n",
|
|
" \"\"\"Unique identifier for the session.\"\"\"\n",
|
|
"\n",
|
|
" created_at: datetime.datetime\n",
|
|
" \"\"\"Timestamp of when the session was created.\"\"\"\n",
|
|
"\n",
|
|
" current_turn: Turn | None\n",
|
|
" \"\"\"The current conversation turn.\"\"\"\n",
|
|
"\n",
|
|
" turns: list[Turn]\n",
|
|
" \"\"\"List of all conversation turns in the session.\"\"\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "435eaf5cf9ef"
|
|
},
|
|
"source": [
|
|
"### Streamed function calling helper functions"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {
|
|
"id": "ec23c165a287"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"async def run_function_async(\n",
|
|
" function: Callable[..., pydantic.BaseModel | Awaitable[pydantic.BaseModel]],\n",
|
|
" function_kwargs: Mapping[str, Any],\n",
|
|
"):\n",
|
|
" \"\"\"Runs a function asynchronously and wraps the results for google-genai FunctionResponse.\n",
|
|
"\n",
|
|
" This function executes a given function asynchronously, handling both synchronous and asynchronous functions. Note: Sync functions are made asynchronous by running in the default threadpool executor so any sync functions should be thread-safe.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" function: The function to execute.\n",
|
|
" function_kwargs: The arguments to pass to the function.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A dictionary containing the function's result or an error message.\n",
|
|
" \"\"\"\n",
|
|
" try:\n",
|
|
" if inspect.iscoroutinefunction(function):\n",
|
|
" fn_result = await function(**function_kwargs)\n",
|
|
" else:\n",
|
|
" loop = asyncio.get_running_loop()\n",
|
|
" fn_result = await loop.run_in_executor(\n",
|
|
" None,\n",
|
|
" lambda kwargs: function(**kwargs),\n",
|
|
" function_kwargs,\n",
|
|
" )\n",
|
|
"\n",
|
|
" return {\"result\": fn_result.model_dump(mode=\"json\")}\n",
|
|
"\n",
|
|
" except Exception as e:\n",
|
|
" return {\"error\": str(e)}\n",
|
|
"\n",
|
|
"\n",
|
|
"async def generate_content_stream(\n",
|
|
" model: str,\n",
|
|
" contents: list[genai_types.Content],\n",
|
|
" config: genai_types.GenerateContentConfig,\n",
|
|
" client: genai.Client,\n",
|
|
" max_recursion_depth: int = 3,\n",
|
|
" fn_map: dict[str, Callable] | None = None,\n",
|
|
") -> AsyncGenerator[genai_types.Content, None]:\n",
|
|
" \"\"\"Streams generated content from a Gemini model, handling function calls within the stream.\n",
|
|
"\n",
|
|
" This function iteratively generates content from a Gemini model, processing function calls\n",
|
|
" encountered during generation. It executes these function calls asynchronously and feeds\n",
|
|
" their results back to the model for continued generation.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" model: The name of the Gemini model to use.\n",
|
|
" contents: The list of Content objects representing the conversation history.\n",
|
|
" config: The GenerateContentConfig for the model.\n",
|
|
" client: The Gemini client.\n",
|
|
" max_recursion_depth: The maximum depth of recursive function calls to prevent infinite loops.\n",
|
|
" **fn_map: A mapping of function names to their corresponding callable functions.\n",
|
|
"\n",
|
|
" Yields:\n",
|
|
" Content objects representing the generated content, including text and function call responses.\n",
|
|
" \"\"\"\n",
|
|
" fn_map = fn_map or {}\n",
|
|
"\n",
|
|
" if max_recursion_depth < 0:\n",
|
|
" print(\"Maximum depth reached, stopping generation.\")\n",
|
|
" return\n",
|
|
"\n",
|
|
" response: AsyncIterator[\n",
|
|
" genai_types.GenerateContentResponse\n",
|
|
" ] = await client.aio.models.generate_content_stream(\n",
|
|
" model=model,\n",
|
|
" contents=contents,\n",
|
|
" config=config,\n",
|
|
" )\n",
|
|
"\n",
|
|
" # iterate over chunk in main request\n",
|
|
" async for chunk in response:\n",
|
|
" if chunk.candidates is None or chunk.candidates[0].content is None:\n",
|
|
" print(\"no candidates or content, skipping chunk.\")\n",
|
|
" continue\n",
|
|
"\n",
|
|
" # yield current chunk content (assume only one candidate)\n",
|
|
" content = chunk.candidates[0].content\n",
|
|
" yield content\n",
|
|
"\n",
|
|
" # if any function calls, execute each in parallel, then call generate after responses are gathered\n",
|
|
" if chunk.function_calls:\n",
|
|
" # create asyncio tasks to execute each function call\n",
|
|
" tasks = list[asyncio.Task[dict[str, Any]]]()\n",
|
|
" for function_call in chunk.function_calls:\n",
|
|
" if function_call.name is None:\n",
|
|
" print(\"skipping function call without name\")\n",
|
|
" continue\n",
|
|
"\n",
|
|
" if function_call.name not in fn_map:\n",
|
|
" raise RuntimeError(\n",
|
|
" f\"Function not provided in fn_map: {function_call.name}\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" func = fn_map[function_call.name]\n",
|
|
" kwargs = function_call.args or {}\n",
|
|
"\n",
|
|
" tasks.append(asyncio.create_task(run_function_async(func, kwargs)))\n",
|
|
"\n",
|
|
" fn_results = await asyncio.gather(*tasks)\n",
|
|
"\n",
|
|
" # create and yield content from function responses\n",
|
|
" fn_response_content = genai_types.Content(\n",
|
|
" role=\"user\",\n",
|
|
" parts=[\n",
|
|
" genai_types.Part.from_function_response(\n",
|
|
" name=fn_call.name, response=fn_result\n",
|
|
" )\n",
|
|
" for fn_call, fn_result in zip(chunk.function_calls, fn_results)\n",
|
|
" ],\n",
|
|
" )\n",
|
|
" yield fn_response_content\n",
|
|
"\n",
|
|
" # continue generation and yield resulting content\n",
|
|
" async for content in generate_content_stream(\n",
|
|
" model=model,\n",
|
|
" contents=contents\n",
|
|
" + [\n",
|
|
" content.model_copy(deep=True),\n",
|
|
" fn_response_content.model_copy(deep=True),\n",
|
|
" ],\n",
|
|
" config=config,\n",
|
|
" client=client,\n",
|
|
" max_recursion_depth=max_recursion_depth - 1,\n",
|
|
" fn_map=fn_map,\n",
|
|
" ):\n",
|
|
" yield content"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "feebcff4fe12"
|
|
},
|
|
"source": [
|
|
"### Tools"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dda88fefd019"
|
|
},
|
|
"source": [
|
|
"#### Find Products Tool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "4a2af10a5ff8"
|
|
},
|
|
"source": [
|
|
"Create a function to construct a BigQuery query to search for products without ranking by product semantic similarity."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {
|
|
"id": "dba908b3486d"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def build_query_without_vector_search(\n",
|
|
" cymbal_products_table_uri: str,\n",
|
|
" cymbal_inventory_table_uri: str,\n",
|
|
" max_results: int = 3,\n",
|
|
" store_ids: list[str] | None = None,\n",
|
|
" min_price: int | None = None,\n",
|
|
" max_price: int | None = None,\n",
|
|
"):\n",
|
|
" \"\"\"Builds a BigQuery SQL query for product search without semantic vector search.\n",
|
|
"\n",
|
|
" This function constructs a SQL query that filters products based on store availability and price range,\n",
|
|
" without using vector search for semantic similarity.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" max_results: The maximum number of results to return.\n",
|
|
" store_ids: Optional list of store IDs to filter products by availability.\n",
|
|
" min_price: Optional minimum price to filter products.\n",
|
|
" max_price: Optional maximum price to filter products.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A tuple containing the SQL query string and the BigQuery query job configuration.\n",
|
|
" \"\"\"\n",
|
|
" where_conditions = list[str]()\n",
|
|
" query_parameters = list[\n",
|
|
" bigquery.ScalarQueryParameter | bigquery.ArrayQueryParameter\n",
|
|
" ]()\n",
|
|
"\n",
|
|
" if min_price is not None:\n",
|
|
" min_price_selector = \"@min_price <= IFNULL(sale_price, list_price)\"\n",
|
|
" where_conditions.append(min_price_selector)\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"min_price\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=min_price,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" if max_price is not None:\n",
|
|
" max_price_selector = \"IFNULL(sale_price, list_price) <= @max_price\"\n",
|
|
" where_conditions.append(max_price_selector)\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"max_price\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=max_price,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" select_products_query = f\"SELECT * FROM `{cymbal_products_table_uri}`\"\n",
|
|
" filter_store_ids = store_ids is not None and len(store_ids) > 0\n",
|
|
" if filter_store_ids:\n",
|
|
" select_products_query = f\"\"\"\n",
|
|
"SELECT\n",
|
|
" product.*\n",
|
|
"FROM\n",
|
|
" `{cymbal_products_table_uri}` AS product,\n",
|
|
" (SELECT DISTINCT\n",
|
|
" uniq_id\n",
|
|
" FROM\n",
|
|
" `{cymbal_inventory_table_uri}`\n",
|
|
" WHERE\n",
|
|
" store_id IN UNNEST(@store_ids)\n",
|
|
" ) AS inventory\n",
|
|
"WHERE\n",
|
|
" product.uniq_id = inventory.uniq_id\n",
|
|
"\"\"\".strip()\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ArrayQueryParameter(\n",
|
|
" name=\"store_ids\",\n",
|
|
" array_type=bigquery.SqlParameterScalarTypes.STRING,\n",
|
|
" values=store_ids,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" where_clause = \"\"\n",
|
|
" if where_conditions:\n",
|
|
" where_clause = f\"WHERE {'AND '.join(f'({cond})' for cond in where_conditions)}\"\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"max_results\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.INTEGER,\n",
|
|
" value=max_results,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" query = f\"\"\"\n",
|
|
"SELECT\n",
|
|
" product.uniq_id AS id,\n",
|
|
" product.product_name AS name,\n",
|
|
" product.product_url AS url,\n",
|
|
" product.product_description AS description,\n",
|
|
" product.brand,\n",
|
|
" product.category,\n",
|
|
" product.available,\n",
|
|
" product.list_price,\n",
|
|
" IF(product.sale_price <= product.list_price, product.sale_price, NULL) AS sale_price\n",
|
|
"FROM\n",
|
|
" ({select_products_query}) AS product\n",
|
|
"{where_clause}\n",
|
|
"LIMIT @max_results\n",
|
|
"\"\"\".strip()\n",
|
|
"\n",
|
|
" query_job_config = bigquery.QueryJobConfig()\n",
|
|
" query_job_config.query_parameters = query_parameters\n",
|
|
"\n",
|
|
" return query, query_job_config"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "5415608141be"
|
|
},
|
|
"source": [
|
|
"Create a function to construct a BigQuery query to search for products and re-rank/filter by product semantic similarity."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {
|
|
"id": "6d84e7983473"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def build_query_with_vector_search(\n",
|
|
" cymbal_products_table_uri: str,\n",
|
|
" cymbal_inventory_table_uri: str,\n",
|
|
" cymbal_embedding_model_uri: str,\n",
|
|
" product_search_query: str,\n",
|
|
" max_results: int = 3,\n",
|
|
" store_ids: list[str] | None = None,\n",
|
|
" min_price: int | None = None,\n",
|
|
" max_price: int | None = None,\n",
|
|
"):\n",
|
|
" \"\"\"Builds a BigQuery SQL query for product search using semantic vector search.\n",
|
|
"\n",
|
|
" This function constructs a SQL query that uses BigQuery's vector search capabilities to find products\n",
|
|
" based on semantic similarity to the provided search query. It also applies filters for store availability\n",
|
|
" and price range.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" product_search_query: The text query for semantic product search.\n",
|
|
" max_results: The maximum number of results to return.\n",
|
|
" store_ids: Optional list of store IDs to filter products by availability.\n",
|
|
" min_price: Optional minimum price to filter products.\n",
|
|
" max_price: Optional maximum price to filter products.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A tuple containing the SQL query string and the BigQuery query job configuration.\n",
|
|
" \"\"\"\n",
|
|
" where_conditions = list[str]()\n",
|
|
" query_parameters = list[\n",
|
|
" bigquery.ScalarQueryParameter | bigquery.ArrayQueryParameter\n",
|
|
" ]()\n",
|
|
"\n",
|
|
" if min_price is not None:\n",
|
|
" min_price_selector = \"@min_price <= IFNULL(base.sale_price, base.list_price)\"\n",
|
|
" where_conditions.append(min_price_selector)\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"min_price\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=min_price,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" if max_price is not None:\n",
|
|
" max_price_selector = \"IFNULL(base.sale_price, base.list_price) <= @max_price\"\n",
|
|
" where_conditions.append(max_price_selector)\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"max_price\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=max_price,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" from_inventory_snippet = \"\"\n",
|
|
" filter_store_ids = store_ids is not None and len(store_ids) > 0\n",
|
|
" if filter_store_ids:\n",
|
|
" where_conditions.append(\"base.uniq_id = inventory.uniq_id\")\n",
|
|
" from_inventory_snippet = f\"\"\",\n",
|
|
"(SELECT DISTINCT\n",
|
|
" uniq_id\n",
|
|
"FROM\n",
|
|
" `{cymbal_inventory_table_uri}`\n",
|
|
"WHERE\n",
|
|
" store_id IN UNNEST(@store_ids)\n",
|
|
") AS inventory\"\"\".strip()\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ArrayQueryParameter(\n",
|
|
" name=\"store_ids\",\n",
|
|
" array_type=bigquery.SqlParameterScalarTypes.STRING,\n",
|
|
" values=store_ids,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" where_clause = \"\"\n",
|
|
" if where_conditions:\n",
|
|
" where_clause = f\"WHERE {'AND '.join(f'({cond})' for cond in where_conditions)}\"\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"top_k\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.INTEGER,\n",
|
|
" value=max_results * 3, # add some wiggle room for post-filtering\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"max_results\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.INTEGER,\n",
|
|
" value=max_results,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" if product_search_query:\n",
|
|
" query = f\"\"\"\n",
|
|
"SELECT\n",
|
|
" base.uniq_id AS id,\n",
|
|
" base.product_name AS name,\n",
|
|
" base.product_url AS url,\n",
|
|
" base.product_description AS description,\n",
|
|
" base.brand,\n",
|
|
" base.category,\n",
|
|
" base.available,\n",
|
|
" base.list_price,\n",
|
|
" IF(base.sale_price <= base.list_price, base.sale_price, NULL) AS sale_price\n",
|
|
"FROM\n",
|
|
" VECTOR_SEARCH(\n",
|
|
" TABLE `{cymbal_products_table_uri}`,\n",
|
|
" 'text_embedding',\n",
|
|
" (SELECT\n",
|
|
" text_embedding,\n",
|
|
" content AS query\n",
|
|
" FROM\n",
|
|
" ML.GENERATE_TEXT_EMBEDDING(\n",
|
|
" MODEL `{cymbal_embedding_model_uri}`,\n",
|
|
" (\n",
|
|
" SELECT @semantic_search_query AS content\n",
|
|
" )\n",
|
|
" )\n",
|
|
" ),\n",
|
|
" top_k => @top_k\n",
|
|
" ) as vector_search\n",
|
|
" {from_inventory_snippet}\n",
|
|
"{where_clause}\n",
|
|
"LIMIT @max_results\n",
|
|
"\"\"\".strip()\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"semantic_search_query\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.STRING,\n",
|
|
" value=product_search_query,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_job_config = bigquery.QueryJobConfig()\n",
|
|
" query_job_config.query_parameters = query_parameters\n",
|
|
"\n",
|
|
" return query, query_job_config"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "d8558f27f50c"
|
|
},
|
|
"source": [
|
|
"Define the function declaration for the product search tool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {
|
|
"id": "ffa8d9816c6e"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def generate_find_products_handler(\n",
|
|
" project: str,\n",
|
|
" cymbal_dataset_location: str,\n",
|
|
" cymbal_products_table_uri: str,\n",
|
|
" cymbal_inventory_table_uri: str,\n",
|
|
" cymbal_embedding_model_uri: str,\n",
|
|
"):\n",
|
|
" \"\"\"Generates a function handler for finding products based on search queries and filters.\n",
|
|
"\n",
|
|
" This function allows searching for products using a semantic search query, filtering by store IDs,\n",
|
|
" and applying price range constraints. It leverages BigQuery's vector search capabilities for\n",
|
|
" semantic similarity and standard SQL queries for other filters.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" project (str): The Google Cloud project ID.\n",
|
|
" cymbal_dataset_location (str): The location of the BigQuery dataset.\n",
|
|
" cymbal_products_table_uri (str): The URI of the products table in BigQuery.\n",
|
|
" cymbal_inventory_table_uri (str): The URI of the inventory table in BigQuery.\n",
|
|
" cymbal_embedding_model_uri (str): The URI of the embedding model in BigQuery for semantic search.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" Callable[[int, Optional[str], Optional[list[str]], Optional[int], Optional[int]], types.ProductSearchResult]:\n",
|
|
" A function that accepts search parameters and returns a ProductSearchResult object.\n",
|
|
"\n",
|
|
" The returned function accepts the following arguments:\n",
|
|
" - max_results (int): The maximum number of products to return.\n",
|
|
" - product_search_query (Optional[str]): A text query for semantic product search.\n",
|
|
" - store_ids (Optional[list[str]]): A list of store IDs to filter products by.\n",
|
|
" - min_price (Optional[int]): The minimum price of the products.\n",
|
|
" - max_price (Optional[int]): The maximum price of the products.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def find_products(\n",
|
|
" max_results: int = 3,\n",
|
|
" product_search_query: str | None = None,\n",
|
|
" store_ids: list[str] | None = None,\n",
|
|
" min_price: int | None = None,\n",
|
|
" max_price: int | None = None,\n",
|
|
" ):\n",
|
|
" \"\"\"Search for products with optional semantic search queries and filters.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" max_results (int): The max number of results to be returned.\n",
|
|
" product_search_query (Optional[str]): Product text search for semantic similarity, can utilize name, description, brand or category.\n",
|
|
" store_ids (Optional[list[str]]): List of store IDs that must carry the returned products. Only include if store IDs are already known, otherwise the store search tool may be more useful.\n",
|
|
" min_price (Optional[int]): Minimum price of products in dollars.\n",
|
|
" max_price (Optional[int]): Maximum price of products in dollars.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" ProductSearchResult: The return value. Object including top matched products and/or an error message.\n",
|
|
" \"\"\"\n",
|
|
" nonlocal \\\n",
|
|
" project, \\\n",
|
|
" cymbal_dataset_location, \\\n",
|
|
" cymbal_products_table_uri, \\\n",
|
|
" cymbal_inventory_table_uri, \\\n",
|
|
" cymbal_embedding_model_uri\n",
|
|
"\n",
|
|
" if max_results >= MAX_PRODUCT_RESULTS:\n",
|
|
" print(\n",
|
|
" f\"Top k is too large ({max_results}). Setting to {MAX_PRODUCT_RESULTS}...\"\n",
|
|
" )\n",
|
|
" max_results = MAX_PRODUCT_RESULTS\n",
|
|
"\n",
|
|
" if product_search_query:\n",
|
|
" query, query_job_config = build_query_with_vector_search(\n",
|
|
" cymbal_products_table_uri=cymbal_products_table_uri,\n",
|
|
" cymbal_inventory_table_uri=cymbal_inventory_table_uri,\n",
|
|
" cymbal_embedding_model_uri=cymbal_embedding_model_uri,\n",
|
|
" max_results=max_results,\n",
|
|
" product_search_query=product_search_query,\n",
|
|
" store_ids=store_ids,\n",
|
|
" min_price=min_price,\n",
|
|
" max_price=max_price,\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" query, query_job_config = build_query_without_vector_search(\n",
|
|
" cymbal_products_table_uri=cymbal_products_table_uri,\n",
|
|
" cymbal_inventory_table_uri=cymbal_inventory_table_uri,\n",
|
|
" max_results=max_results,\n",
|
|
" store_ids=store_ids,\n",
|
|
" min_price=min_price,\n",
|
|
" max_price=max_price,\n",
|
|
" )\n",
|
|
"\n",
|
|
" bq_client = bigquery.Client(project=project, location=cymbal_dataset_location)\n",
|
|
" query_job = bq_client.query(\n",
|
|
" query=query,\n",
|
|
" job_config=query_job_config,\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_df = query_job.to_dataframe()\n",
|
|
"\n",
|
|
" products = [\n",
|
|
" Product.model_validate(row.to_dict()) for idx, row in query_df.iterrows()\n",
|
|
" ]\n",
|
|
"\n",
|
|
" return ProductSearchResult(products=products, query=query)\n",
|
|
"\n",
|
|
" return find_products\n",
|
|
"\n",
|
|
"\n",
|
|
"find_products_fd = genai_types.FunctionDeclaration(\n",
|
|
" response=None,\n",
|
|
" description=\"Search for products with optional semantic search queries and filters.\",\n",
|
|
" name=\"find_products\",\n",
|
|
" parameters=genai_types.Schema(\n",
|
|
" properties={\n",
|
|
" \"max_results\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" default=3,\n",
|
|
" description=\"The max number of results to be returned.\",\n",
|
|
" ),\n",
|
|
" \"product_search_query\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.STRING,\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"Product text search for semantic similarity, can utilize name, description, brand or category.\",\n",
|
|
" ),\n",
|
|
" \"store_ids\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.ARRAY,\n",
|
|
" items=genai_types.Schema(type=genai_types.Type.STRING),\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"List of store IDs that must carry the returned products. Only include if store IDs are already known, otherwise the store search tool may be more useful.\",\n",
|
|
" ),\n",
|
|
" \"min_price\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"Minimum price of products in dollars\",\n",
|
|
" ),\n",
|
|
" \"max_price\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"Maximum price of products in dollars\",\n",
|
|
" ),\n",
|
|
" },\n",
|
|
" required=[],\n",
|
|
" type=genai_types.Type.OBJECT,\n",
|
|
" ),\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "05d2dc49ca95"
|
|
},
|
|
"source": [
|
|
"Test find_products tool function"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"metadata": {
|
|
"id": "63c79562fac6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"find_products = generate_find_products_handler(\n",
|
|
" project=PROJECT_ID,\n",
|
|
" cymbal_dataset_location=CYMBAL_DATASET_LOCATION,\n",
|
|
" cymbal_products_table_uri=CYMBAL_PRODUCTS_TABLE_URI,\n",
|
|
" cymbal_inventory_table_uri=CYMBAL_INVENTORY_TABLE_URI,\n",
|
|
" cymbal_embedding_model_uri=CYMBAL_EMBEDDING_MODEL_URI,\n",
|
|
")\n",
|
|
"\n",
|
|
"product_search_result = find_products(\n",
|
|
" max_results=5,\n",
|
|
" product_search_query=\"curtains\",\n",
|
|
" max_price=100.0,\n",
|
|
")\n",
|
|
"[(product.id, product.name) for product in product_search_result.products]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "0d2295f77807"
|
|
},
|
|
"source": [
|
|
"#### Find Nearby Stores Tool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"metadata": {
|
|
"id": "9090613690c7"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"MAX_STORE_RESULTS = 10\n",
|
|
"STORE_NAME_SIMILARITY_THRESHOLD = 90\n",
|
|
"\n",
|
|
"\n",
|
|
"def generate_find_stores_handler(\n",
|
|
" project: str,\n",
|
|
" cymbal_dataset_location: str,\n",
|
|
" cymbal_stores_table_uri: str,\n",
|
|
" cymbal_inventory_table_uri: str,\n",
|
|
" user_latitude: float | None = None,\n",
|
|
" user_longitude: float | None = None,\n",
|
|
"):\n",
|
|
" \"\"\"Generates a handler function for finding stores based on various criteria.\n",
|
|
"\n",
|
|
" This function creates a closure that encapsulates user location information (latitude and longitude)\n",
|
|
" and returns a function that can search for stores. The returned function can filter stores based on\n",
|
|
" product IDs, proximity (radius), store name, and maximum number of results.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" project (str): The Google Cloud project ID.\n",
|
|
" cymbal_dataset_location (str): The location of the BigQuery dataset.\n",
|
|
" cymbal_stores_table_uri (str): The URI of the BigQuery table containing store information.\n",
|
|
" cymbal_inventory_table_uri (str): The URI of the BigQuery table containing product inventory.\n",
|
|
" user_latitude (Optional[float]): The user's latitude for location-based searches.\n",
|
|
" user_longitude (Optional[float]): The user's longitude for location-based searches.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" Callable: A function that takes product IDs, max results, radius, and store name as input and returns a StoreSearchResult.\n",
|
|
"\n",
|
|
" Raises:\n",
|
|
" AssertionError: If only one of `user_latitude` or `user_longitude` is provided.\n",
|
|
" \"\"\"\n",
|
|
" assert not ((user_latitude is None) ^ (user_longitude is None)), (\n",
|
|
" \"Latitude and longitude must both be defined or both null\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" def find_stores(\n",
|
|
" product_ids: list[str] | None = None,\n",
|
|
" max_results: int = 3,\n",
|
|
" # Note: google-genai doesn't properly handle floats, so we just set this as an integer\n",
|
|
" radius_km: int | None = None,\n",
|
|
" store_name: str | None = None,\n",
|
|
" ):\n",
|
|
" \"\"\"Search for stores nearby, by name, or offering certain products.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" max_results (int): The max number of results to be returned. Largest allowed value is 10.\n",
|
|
" radius_km (Optional[int]): Radius in kilometers to restrict the nearby search around the user location. The user location doesn't have to be provided in the conversation context. This function can retrieve the user location from a backend database.\n",
|
|
" store_name (Optional[str]): The name (or part of a name) of a store to search for. Will try to find stores that fuzzy match this name.\n",
|
|
" product_ids (list[str]): List of product IDs that must exist at the given store. Leave empty if there is no product ID filtering.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" StoreSearchResult: The return value. Object including top matched stores and/or an error message.\n",
|
|
" \"\"\"\n",
|
|
" nonlocal \\\n",
|
|
" project, \\\n",
|
|
" cymbal_dataset_location, \\\n",
|
|
" cymbal_stores_table_uri, \\\n",
|
|
" cymbal_inventory_table_uri, \\\n",
|
|
" user_latitude, \\\n",
|
|
" user_longitude\n",
|
|
"\n",
|
|
" product_ids = product_ids or []\n",
|
|
"\n",
|
|
" query_parameters = list[\n",
|
|
" bigquery.ScalarQueryParameter | bigquery.ArrayQueryParameter\n",
|
|
" ]()\n",
|
|
"\n",
|
|
" if max_results >= MAX_STORE_RESULTS:\n",
|
|
" print(\n",
|
|
" f\"Top k is too large ({max_results}). Setting to {MAX_STORE_RESULTS}...\"\n",
|
|
" )\n",
|
|
" max_results = MAX_STORE_RESULTS\n",
|
|
"\n",
|
|
" radius_selector = None\n",
|
|
" if radius_km:\n",
|
|
" if user_latitude is None or user_longitude is None:\n",
|
|
" raise ValueError(\"User location is not known\")\n",
|
|
"\n",
|
|
" radius_selector = \"ST_DISTANCE(ST_GEOGPOINT(@longitude, @latitude), ST_GEOGPOINT(longitude, latitude)) <= @radius_meters\"\n",
|
|
"\n",
|
|
" query_parameters.extend(\n",
|
|
" [\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"latitude\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=user_latitude,\n",
|
|
" ),\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"longitude\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=user_longitude,\n",
|
|
" ),\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"radius_meters\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.FLOAT,\n",
|
|
" value=radius_km * 1_000.0,\n",
|
|
" ),\n",
|
|
" ]\n",
|
|
" )\n",
|
|
"\n",
|
|
" where_clause = \"\"\n",
|
|
" if radius_selector:\n",
|
|
" where_clause = f\"WHERE {radius_selector}\"\n",
|
|
"\n",
|
|
" select_stores_query = f\"SELECT * FROM {cymbal_stores_table_uri}\"\n",
|
|
" if len(product_ids) > 0:\n",
|
|
" select_stores_query = f\"\"\"\n",
|
|
" SELECT\n",
|
|
" store.*\n",
|
|
" FROM\n",
|
|
" {cymbal_stores_table_uri} AS store,\n",
|
|
" (SELECT DISTINCT\n",
|
|
" store_id\n",
|
|
" FROM\n",
|
|
" {cymbal_inventory_table_uri}\n",
|
|
" WHERE\n",
|
|
" uniq_id IN UNNEST(@product_ids)\n",
|
|
" ) AS inventory\n",
|
|
" WHERE\n",
|
|
" store.store_id = inventory.store_id\n",
|
|
" \"\"\".strip()\n",
|
|
"\n",
|
|
" query_parameters.append(\n",
|
|
" bigquery.ArrayQueryParameter(\n",
|
|
" name=\"product_ids\",\n",
|
|
" array_type=bigquery.SqlParameterScalarTypes.STRING,\n",
|
|
" values=product_ids,\n",
|
|
" )\n",
|
|
" )\n",
|
|
"\n",
|
|
" query = f\"\"\"\n",
|
|
" SELECT\n",
|
|
" CAST(store.store_id AS STRING) AS id,\n",
|
|
" store.name,\n",
|
|
" store.url,\n",
|
|
" store.street_address,\n",
|
|
" store.city,\n",
|
|
" store.state,\n",
|
|
" store.zip_code,\n",
|
|
" store.country,\n",
|
|
" store.phone_number_1 as phone_number,\n",
|
|
" store.latitude,\n",
|
|
" store.longitude,\n",
|
|
" FROM\n",
|
|
" ({select_stores_query}) AS store\n",
|
|
" {where_clause}\n",
|
|
" \"\"\".strip()\n",
|
|
"\n",
|
|
" query_job_config = bigquery.QueryJobConfig()\n",
|
|
" query_job_config.query_parameters = query_parameters\n",
|
|
"\n",
|
|
" bq_client = bigquery.Client(project=project, location=cymbal_dataset_location)\n",
|
|
" query_job = bq_client.query(\n",
|
|
" query=query,\n",
|
|
" job_config=query_job_config,\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_df = query_job.to_dataframe()\n",
|
|
"\n",
|
|
" stores = [\n",
|
|
" Store.model_validate(row.to_dict())\n",
|
|
" for idx, row in query_df.iterrows()\n",
|
|
" # Filter by store names fuzzy matching the user input\n",
|
|
" if (\n",
|
|
" store_name is None\n",
|
|
" or fuzz.partial_ratio(row[\"name\"], store_name)\n",
|
|
" >= STORE_NAME_SIMILARITY_THRESHOLD\n",
|
|
" )\n",
|
|
" ][:max_results] # filter max results\n",
|
|
"\n",
|
|
" return StoreSearchResult(stores=stores, query=query)\n",
|
|
"\n",
|
|
" return find_stores\n",
|
|
"\n",
|
|
"\n",
|
|
"find_stores_fd = genai_types.FunctionDeclaration(\n",
|
|
" response=None,\n",
|
|
" description=\"Search for stores nearby, by name, or offering certain products.\",\n",
|
|
" name=\"find_stores\",\n",
|
|
" parameters=genai_types.Schema(\n",
|
|
" properties={\n",
|
|
" \"max_results\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" default=3,\n",
|
|
" description=\"The max number of results to be returned.\",\n",
|
|
" ),\n",
|
|
" \"store_name\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.STRING,\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"The name (or part of a name) of a store to search for. Will try to find stores that fuzzy match this name.\",\n",
|
|
" ),\n",
|
|
" \"product_ids\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.ARRAY,\n",
|
|
" items=genai_types.Schema(type=genai_types.Type.STRING),\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"List of product IDs that must exist at the given store. Leave empty if there is no product ID filtering.\",\n",
|
|
" ),\n",
|
|
" \"radius_km\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" nullable=True,\n",
|
|
" default=None,\n",
|
|
" description=\"Radius in kilometers to restrict the nearby search around the user location. The user location doesn't have to be provided in the conversation context. This function can retrieve the user location from a backend database.\",\n",
|
|
" ),\n",
|
|
" },\n",
|
|
" required=[],\n",
|
|
" type=genai_types.Type.OBJECT,\n",
|
|
" ),\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "83665f5bbc88"
|
|
},
|
|
"source": [
|
|
"test find_stores tool function by radius"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 10,
|
|
"metadata": {
|
|
"id": "55587c98a4ac"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"user_latitude = 32.631012\n",
|
|
"user_longitude = -116.968043\n",
|
|
"\n",
|
|
"find_stores = generate_find_stores_handler(\n",
|
|
" project=PROJECT_ID,\n",
|
|
" cymbal_dataset_location=CYMBAL_DATASET_LOCATION,\n",
|
|
" cymbal_stores_table_uri=CYMBAL_STORES_TABLE_URI,\n",
|
|
" cymbal_inventory_table_uri=CYMBAL_INVENTORY_TABLE_URI,\n",
|
|
" user_latitude=user_latitude,\n",
|
|
" user_longitude=user_longitude,\n",
|
|
")\n",
|
|
"\n",
|
|
"store_search_result = find_stores(\n",
|
|
" product_ids=[],\n",
|
|
" max_results=5,\n",
|
|
" radius_km=50,\n",
|
|
")\n",
|
|
"\n",
|
|
"[(store.id, store.name) for store in store_search_result.stores]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "fb7b3e9b470d"
|
|
},
|
|
"source": [
|
|
"test find_stores tool function by product"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"metadata": {
|
|
"id": "c2e646491b9c"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"user_latitude = 32.631012\n",
|
|
"user_longitude = -116.968043\n",
|
|
"\n",
|
|
"find_stores = generate_find_stores_handler(\n",
|
|
" project=PROJECT_ID,\n",
|
|
" cymbal_dataset_location=CYMBAL_DATASET_LOCATION,\n",
|
|
" cymbal_stores_table_uri=CYMBAL_STORES_TABLE_URI,\n",
|
|
" cymbal_inventory_table_uri=CYMBAL_INVENTORY_TABLE_URI,\n",
|
|
" user_latitude=user_latitude,\n",
|
|
" user_longitude=user_longitude,\n",
|
|
")\n",
|
|
"\n",
|
|
"# select random product\n",
|
|
"target_product = random.choice(product_search_result.products)\n",
|
|
"\n",
|
|
"store_search_result = find_stores(\n",
|
|
" product_ids=[target_product.id],\n",
|
|
" max_results=5,\n",
|
|
")\n",
|
|
"\n",
|
|
"[(store.id, store.name) for store in store_search_result.stores]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dacaadac3c8b"
|
|
},
|
|
"source": [
|
|
"#### Find Inventory Tool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"metadata": {
|
|
"id": "c86037aa0f7b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def generate_find_inventory_handler(\n",
|
|
" project: str,\n",
|
|
" cymbal_inventory_table_uri: str,\n",
|
|
" cymbal_dataset_location: str,\n",
|
|
"):\n",
|
|
" \"\"\"Generates a handler function for finding inventory information.\n",
|
|
"\n",
|
|
" This function creates a callable that queries BigQuery to find inventory\n",
|
|
" details for a specific product at a given store.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" project (str): The Google Cloud project ID.\n",
|
|
" cymbal_inventory_table_uri (str): The URI of the BigQuery table containing inventory information.\n",
|
|
" cymbal_dataset_location (str): The location of the BigQuery dataset.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" Callable: A function that takes store ID and product ID as input and returns an InventorySearchResult.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def find_inventory(\n",
|
|
" store_id: int,\n",
|
|
" product_id: str,\n",
|
|
" ):\n",
|
|
" \"\"\"Look up the inventory query for a given product at a certain store. The product ID and store ID must be known before calling this function. If either are not known, use the other tools to first find the right store and product IDs.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" store_id (int): Unique identifier of the store.\n",
|
|
" product_id (str): Unique identifier of the product.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" InventorySearchResult: The return value. Object including the current inventory and/or an error message.\n",
|
|
" \"\"\"\n",
|
|
" nonlocal project, cymbal_inventory_table_uri\n",
|
|
"\n",
|
|
" query = f\"\"\"\n",
|
|
" SELECT\n",
|
|
" CAST(store_id AS STRING) AS store_id,\n",
|
|
" uniq_id AS product_id,\n",
|
|
" inventory AS value,\n",
|
|
" FROM\n",
|
|
" {cymbal_inventory_table_uri}\n",
|
|
" WHERE\n",
|
|
" store_id = @store_id\n",
|
|
" AND uniq_id = @product_id\n",
|
|
" \"\"\".strip()\n",
|
|
"\n",
|
|
" query_job_config = bigquery.QueryJobConfig()\n",
|
|
" query_job_config.query_parameters = [\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"store_id\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.INTEGER,\n",
|
|
" value=store_id,\n",
|
|
" ),\n",
|
|
" bigquery.ScalarQueryParameter(\n",
|
|
" name=\"product_id\",\n",
|
|
" type_=bigquery.SqlParameterScalarTypes.STRING,\n",
|
|
" value=product_id,\n",
|
|
" ),\n",
|
|
" ]\n",
|
|
"\n",
|
|
" bq_client = bigquery.Client(project=project, location=cymbal_dataset_location)\n",
|
|
" query_job = bq_client.query(\n",
|
|
" query=query,\n",
|
|
" job_config=query_job_config,\n",
|
|
" )\n",
|
|
"\n",
|
|
" query_df = query_job.to_dataframe()\n",
|
|
"\n",
|
|
" if len(query_df) > 1:\n",
|
|
" return InventorySearchResult(\n",
|
|
" error=\"Multiple store/product combinations found.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" if len(query_df) == 0:\n",
|
|
" return InventorySearchResult(error=\"No store/product combinations found.\")\n",
|
|
"\n",
|
|
" inventory = Inventory.model_validate(query_df.iloc[0].to_dict())\n",
|
|
"\n",
|
|
" return InventorySearchResult(inventory=inventory, query=query)\n",
|
|
"\n",
|
|
" return find_inventory\n",
|
|
"\n",
|
|
"\n",
|
|
"find_inventory_fd = genai_types.FunctionDeclaration(\n",
|
|
" response=None,\n",
|
|
" description=\"Look up the inventory query for a given product at a certain store. The product ID and store ID must be known before calling this function. If either are not known, use the other tools to first find the right store and product IDs.\",\n",
|
|
" name=\"find_inventory\",\n",
|
|
" parameters=genai_types.Schema(\n",
|
|
" properties={\n",
|
|
" \"store_id\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.INTEGER,\n",
|
|
" description=\"Unique identifier of the store.\",\n",
|
|
" ),\n",
|
|
" \"product_id\": genai_types.Schema(\n",
|
|
" type=genai_types.Type.STRING,\n",
|
|
" description=\"Unique identifier of the product.\",\n",
|
|
" ),\n",
|
|
" },\n",
|
|
" required=[\"store_id\", \"product_id\"],\n",
|
|
" type=genai_types.Type.OBJECT,\n",
|
|
" ),\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 13,
|
|
"metadata": {
|
|
"id": "9c17d44e7c82"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"find_inventory = generate_find_inventory_handler(\n",
|
|
" project=PROJECT_ID,\n",
|
|
" cymbal_inventory_table_uri=CYMBAL_INVENTORY_TABLE_URI,\n",
|
|
" cymbal_dataset_location=CYMBAL_DATASET_LOCATION,\n",
|
|
")\n",
|
|
"\n",
|
|
"target_store = random.choice(store_search_result.stores)\n",
|
|
"\n",
|
|
"inventory_search_result = find_inventory(\n",
|
|
" store_id=target_store.id,\n",
|
|
" product_id=target_product.id,\n",
|
|
")\n",
|
|
"inventory_search_result"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "fc3f58a2f47d"
|
|
},
|
|
"source": [
|
|
"### LangGraph Nodes"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "becf44d0f6df"
|
|
},
|
|
"source": [
|
|
"#### Chat Node\n",
|
|
"\n",
|
|
"Generate a response for the user input. If needed, will call and stream the results from available tools:\n",
|
|
"* find_products: Search for products given some stores, price range, and/or open text search query\n",
|
|
"* find_stores: Search for stores given its name, some offered products, and/or a search radius near the user.\n",
|
|
"* find_inventory: Search for the inventory of a given product/store pair."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"metadata": {
|
|
"id": "e05125b83787"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"CHAT_SYSTEM_PROMPT = \"\"\"\"\n",
|
|
"You are a chat assistant for the Cymbal Retail site which manages inventory for stores/businesses across many industries.\n",
|
|
"Help answer any user questions.\n",
|
|
"Whenever you don't know something, use one or more of your tools before responding to retrieve live data about stores, products, and inventory.\n",
|
|
"The purpose of each tool is:\n",
|
|
"- find_products: Search for products given some stores, price range, and/or open text search query\n",
|
|
"- find_stores: Search for stores given its name, some offered products, and/or a search radius near the user.\n",
|
|
"- find_inventory: Search for the inventory of a given product/store pair.\n",
|
|
"Note: the user's location is stored in the persistent session storage. It can be retrieved in the background for the store search radius.\n",
|
|
"\"\"\".strip()\n",
|
|
"\n",
|
|
"\n",
|
|
"async def ainvoke_chat(\n",
|
|
" state: GraphSession,\n",
|
|
" config: lc_config.RunnableConfig,\n",
|
|
") -> lg_types.Command[PostProcessNodeTargetLiteral]:\n",
|
|
" \"\"\"Asynchronously invokes the chat node to generate a response using a Gemini model.\n",
|
|
"\n",
|
|
" This function takes the current conversation state, including the user's input and conversation history,\n",
|
|
" and generates a response using a Gemini model. It supports function calling to retrieve live data about\n",
|
|
" stores, products, and inventory. It streams the response text and updates the conversation state with the\n",
|
|
" generated response and function call results.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" state: The current state of the conversation session.\n",
|
|
" config: The LangChain RunnableConfig (unused in this function).\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A Command object specifying the next node to transition to (post-processing)\n",
|
|
" and the updated conversation state.\n",
|
|
" \"\"\"\n",
|
|
" agent_config = AgentConfig.model_validate(\n",
|
|
" config[\"configurable\"].get(\"agent_config\", {})\n",
|
|
" )\n",
|
|
"\n",
|
|
" stream_writer = get_stream_writer()\n",
|
|
"\n",
|
|
" current_turn = state.get(\"current_turn\")\n",
|
|
" assert current_turn is not None, \"current turn must be set\"\n",
|
|
"\n",
|
|
" user_input = current_turn.get(\"user_input\")\n",
|
|
" assert user_input is not None, \"user input must be set\"\n",
|
|
"\n",
|
|
" user_latitude = current_turn.get(\"user_latitude\")\n",
|
|
" user_longitude = current_turn.get(\"user_longitude\")\n",
|
|
"\n",
|
|
" # Initialize generate model\n",
|
|
" client = genai.Client(\n",
|
|
" vertexai=True,\n",
|
|
" project=agent_config.project,\n",
|
|
" location=agent_config.region,\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Add new user input to history\n",
|
|
" turns = state.get(\"turns\", [])\n",
|
|
" history = [content for turn in turns for content in turn.get(\"messages\", [])]\n",
|
|
" user_content = genai_types.Content(\n",
|
|
" role=\"user\",\n",
|
|
" parts=[genai_types.Part.from_text(text=user_input)],\n",
|
|
" )\n",
|
|
" contents = history + [user_content]\n",
|
|
"\n",
|
|
" response_text = \"\"\n",
|
|
" new_contents = [user_content.model_copy(deep=True)]\n",
|
|
" try:\n",
|
|
" find_stores_handler = generate_find_stores_handler(\n",
|
|
" project=agent_config.project,\n",
|
|
" cymbal_dataset_location=agent_config.cymbal_dataset_location,\n",
|
|
" cymbal_stores_table_uri=agent_config.cymbal_stores_table_uri,\n",
|
|
" cymbal_inventory_table_uri=agent_config.cymbal_inventory_table_uri,\n",
|
|
" user_latitude=user_latitude,\n",
|
|
" user_longitude=user_longitude,\n",
|
|
" )\n",
|
|
" find_products_handler = generate_find_products_handler(\n",
|
|
" project=agent_config.project,\n",
|
|
" cymbal_dataset_location=agent_config.cymbal_dataset_location,\n",
|
|
" cymbal_products_table_uri=agent_config.cymbal_products_table_uri,\n",
|
|
" cymbal_inventory_table_uri=agent_config.cymbal_inventory_table_uri,\n",
|
|
" cymbal_embedding_model_uri=agent_config.cymbal_embedding_model_uri,\n",
|
|
" )\n",
|
|
" find_inventory_handler = generate_find_inventory_handler(\n",
|
|
" project=agent_config.project,\n",
|
|
" cymbal_dataset_location=agent_config.cymbal_dataset_location,\n",
|
|
" cymbal_inventory_table_uri=agent_config.cymbal_inventory_table_uri,\n",
|
|
" )\n",
|
|
"\n",
|
|
" # generate streaming response\n",
|
|
" response = generate_content_stream(\n",
|
|
" model=agent_config.chat_model_name,\n",
|
|
" contents=contents,\n",
|
|
" config=genai_types.GenerateContentConfig(\n",
|
|
" system_instruction=CHAT_SYSTEM_PROMPT,\n",
|
|
" tools=[\n",
|
|
" genai_types.Tool(\n",
|
|
" function_declarations=[\n",
|
|
" find_products_fd,\n",
|
|
" find_stores_fd,\n",
|
|
" find_inventory_fd,\n",
|
|
" ]\n",
|
|
" )\n",
|
|
" ],\n",
|
|
" temperature=0.6,\n",
|
|
" candidate_count=1,\n",
|
|
" seed=42,\n",
|
|
" automatic_function_calling=genai_types.AutomaticFunctionCallingConfig(\n",
|
|
" disable=True\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
" client=client,\n",
|
|
" fn_map={\n",
|
|
" find_products_fd.name: find_products_handler,\n",
|
|
" find_stores_fd.name: find_stores_handler,\n",
|
|
" find_inventory_fd.name: find_inventory_handler,\n",
|
|
" },\n",
|
|
" )\n",
|
|
"\n",
|
|
" async for content in response:\n",
|
|
" used_content = False\n",
|
|
" for part in content.parts:\n",
|
|
" if part.text:\n",
|
|
" response_text += part.text\n",
|
|
" used_content = True\n",
|
|
" stream_writer({\"text\": part.text})\n",
|
|
" if part.function_call:\n",
|
|
" used_content = True\n",
|
|
" stream_writer(\n",
|
|
" {\"function_call\": part.function_call.model_dump(mode=\"json\")}\n",
|
|
" )\n",
|
|
" if part.function_response:\n",
|
|
" used_content = True\n",
|
|
" stream_writer(\n",
|
|
" {\n",
|
|
" \"function_response\": part.function_response.model_dump(\n",
|
|
" mode=\"json\"\n",
|
|
" )\n",
|
|
" }\n",
|
|
" )\n",
|
|
"\n",
|
|
" if used_content:\n",
|
|
" new_contents.append(content.model_copy(deep=True))\n",
|
|
"\n",
|
|
" except Exception as e:\n",
|
|
" logger.exception(e)\n",
|
|
" # unexpected error, display it\n",
|
|
" response_text = f\"An unexpected error occurred during generation, please try again.\\n\\nError = {e!s}\"\n",
|
|
" stream_writer({\"error\": response_text})\n",
|
|
"\n",
|
|
" current_turn[\"response\"] = response_text.strip()\n",
|
|
" current_turn[\"messages\"] = new_contents\n",
|
|
"\n",
|
|
" return lg_types.Command(\n",
|
|
" update=GraphSession(current_turn=current_turn),\n",
|
|
" goto=POST_PROCESS_NODE_NAME,\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "4e2d40f06762"
|
|
},
|
|
"source": [
|
|
"#### Post-Process Node\n",
|
|
"\n",
|
|
"Add current turn to the history and reset current turn."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"metadata": {
|
|
"id": "3a9e69592a6b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"async def ainvoke_post_process(\n",
|
|
" state: GraphSession,\n",
|
|
" config: lc_config.RunnableConfig,\n",
|
|
") -> lg_types.Command[EndNodeTargetLiteral]:\n",
|
|
" \"\"\"Asynchronously invokes the post-processing node to finalize the current conversation turn.\n",
|
|
"\n",
|
|
" This function takes the current conversation state, validates that the current turn and its response are set,\n",
|
|
" adds the completed turn to the conversation history, and resets the current turn. This effectively concludes\n",
|
|
" the processing of the current user input and prepares the session for the next input.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" state: The current state of the conversation session.\n",
|
|
" config: The LangChain RunnableConfig (unused in this function).\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" A Command object specifying the end of the graph execution and the updated conversation state.\n",
|
|
" \"\"\"\n",
|
|
" del config # unused\n",
|
|
"\n",
|
|
" current_turn = state.get(\"current_turn\")\n",
|
|
"\n",
|
|
" assert current_turn is not None, \"Current turn must be set.\"\n",
|
|
" assert current_turn[\"response\"] is not None, (\n",
|
|
" \"Response from current turn must be set.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" turns = state.get(\"turns\", []) + [current_turn]\n",
|
|
"\n",
|
|
" return lg_types.Command(update=GraphSession(current_turn=None, turns=turns))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "8e17c18a7302"
|
|
},
|
|
"source": [
|
|
"## Compile Function Calling Agent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {
|
|
"id": "66a1564df060"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def load_graph():\n",
|
|
" \"\"\"Load state graph for basic example.\"\"\"\n",
|
|
" # Graph\n",
|
|
" state_graph = graph.StateGraph(state_schema=GraphSession)\n",
|
|
"\n",
|
|
" # Nodes\n",
|
|
" state_graph.add_node(CHAT_NODE_NAME, ainvoke_chat)\n",
|
|
" state_graph.add_node(POST_PROCESS_NODE_NAME, ainvoke_post_process)\n",
|
|
" state_graph.set_entry_point(CHAT_NODE_NAME)\n",
|
|
"\n",
|
|
" return state_graph\n",
|
|
"\n",
|
|
"\n",
|
|
"state_graph = load_graph()\n",
|
|
"compiled_graph = state_graph.compile(memory_checkpoint.MemorySaver())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "bb31fa3462d3"
|
|
},
|
|
"source": [
|
|
"### Visualize agent graph"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"metadata": {
|
|
"id": "fd2cb303478f"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"display(ipd.Image(state_graph.compile().get_graph().draw_mermaid_png()))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "96fe80e24098"
|
|
},
|
|
"source": [
|
|
"### Wrapper function to stream generation output to notebook"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 18,
|
|
"metadata": {
|
|
"id": "3b86271fc7f6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"async def ask(\n",
|
|
" user_input: str,\n",
|
|
" session: str | None = None,\n",
|
|
" user_latitude: float = 32.631012,\n",
|
|
" user_longitude: float = -116.968043,\n",
|
|
"):\n",
|
|
" thread_id = session or uuid.uuid4().hex\n",
|
|
"\n",
|
|
" agent_config = AgentConfig(\n",
|
|
" project=PROJECT_ID,\n",
|
|
" region=REGION,\n",
|
|
" chat_model_name=CHAT_MODEL_NAME,\n",
|
|
" cymbal_dataset_location=CYMBAL_DATASET_LOCATION,\n",
|
|
" cymbal_products_table_uri=CYMBAL_PRODUCTS_TABLE_URI,\n",
|
|
" cymbal_stores_table_uri=CYMBAL_STORES_TABLE_URI,\n",
|
|
" cymbal_inventory_table_uri=CYMBAL_INVENTORY_TABLE_URI,\n",
|
|
" cymbal_embedding_model_uri=CYMBAL_EMBEDDING_MODEL_URI,\n",
|
|
" )\n",
|
|
"\n",
|
|
" current_source = last_source = None\n",
|
|
" all_text = \"\"\n",
|
|
" async for chunk in compiled_graph.astream(\n",
|
|
" input={\n",
|
|
" \"current_turn\": {\n",
|
|
" \"user_input\": user_input,\n",
|
|
" \"user_latitude\": user_latitude,\n",
|
|
" \"user_longitude\": user_longitude,\n",
|
|
" }\n",
|
|
" },\n",
|
|
" config={\"configurable\": {\"thread_id\": thread_id, \"agent_config\": agent_config}},\n",
|
|
" stream_mode=\"custom\",\n",
|
|
" ):\n",
|
|
" assert isinstance(chunk, dict), \"Expected dictionary chunk\"\n",
|
|
"\n",
|
|
" text = \"\"\n",
|
|
"\n",
|
|
" if \"text\" in chunk:\n",
|
|
" text = chunk[\"text\"]\n",
|
|
" current_source = \"text\"\n",
|
|
"\n",
|
|
" elif \"function_call\" in chunk:\n",
|
|
" function_call_dict = chunk[\"function_call\"]\n",
|
|
"\n",
|
|
" fn_name = function_call_dict.get(\"name\") or \"unknown\"\n",
|
|
" fn_args = function_call_dict.get(\"args\") or {}\n",
|
|
"\n",
|
|
" fn_args_string = \", \".join(f\"{k}={v}\" for k, v in fn_args.items())\n",
|
|
" fn_string = f\"**{fn_name}**({fn_args_string})\"\n",
|
|
"\n",
|
|
" text = f\"Calling function... {fn_string}\"\n",
|
|
" current_source = \"function_call\"\n",
|
|
"\n",
|
|
" elif \"function_response\" in chunk:\n",
|
|
" function_response_dict = chunk[\"function_response\"]\n",
|
|
"\n",
|
|
" fn_name = function_response_dict.get(\"name\") or \"unknown\"\n",
|
|
"\n",
|
|
" if function_response_dict.get(\"response\") is None:\n",
|
|
" text = f\"Received empty function response (name={fn_name}).\"\n",
|
|
"\n",
|
|
" elif \"result\" in function_response_dict.get(\"response\"):\n",
|
|
" fn_result = function_response_dict[\"response\"][\"result\"]\n",
|
|
" text = \"\\n\\n\".join(\n",
|
|
" [\n",
|
|
" f\"Function result for **{fn_name}**...\",\n",
|
|
" \"```json\",\n",
|
|
" json.dumps(fn_result, indent=2),\n",
|
|
" \"```\",\n",
|
|
" ]\n",
|
|
" )\n",
|
|
"\n",
|
|
" elif \"error\" in function_response_dict.get(\"response\"):\n",
|
|
" fn_result = function_response_dict[\"response\"][\"error\"]\n",
|
|
" text = f\"Function error (name={fn_name})... {fn_result}\"\n",
|
|
"\n",
|
|
" current_source = \"function_response\"\n",
|
|
"\n",
|
|
" elif \"error\" in chunk:\n",
|
|
" text = chunk[\"error\"]\n",
|
|
" current_source = \"error\"\n",
|
|
"\n",
|
|
" else:\n",
|
|
" print(\"unhandled chunk case:\", chunk)\n",
|
|
"\n",
|
|
" if last_source is not None and last_source != current_source:\n",
|
|
" text = \"\\n\\n---\\n\\n\" + text\n",
|
|
"\n",
|
|
" last_source = current_source\n",
|
|
"\n",
|
|
" all_text += text\n",
|
|
" display(ipd.Markdown(all_text), clear=True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "9c85d0110e4d"
|
|
},
|
|
"source": [
|
|
"## Test Conversation"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 19,
|
|
"metadata": {
|
|
"id": "d66bbc9165e7"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"session = uuid.uuid4().hex"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 20,
|
|
"metadata": {
|
|
"id": "b893a03e5d6b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"await ask(\"Do you sell curtains under $50?\", session=session)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 21,
|
|
"metadata": {
|
|
"id": "8bdeff970d29"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"await ask(\n",
|
|
" \"What stores carry the last one?\",\n",
|
|
" session=session,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 22,
|
|
"metadata": {
|
|
"id": "3b8c3ce67081"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"await ask(\n",
|
|
" \"How many do they have left in Chula Vista? That's the closest to me\",\n",
|
|
" session=session,\n",
|
|
")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "function-calling.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|