|
16 | 16 |
|
17 | 17 | from google.adk.agents.llm_agent import LlmAgent |
18 | 18 | from google.adk.auth.auth_credential import AuthCredentialTypes |
| 19 | +from google.adk.tools.bigtable import query_tool as bigtable_query_tool |
19 | 20 | from google.adk.tools.bigtable.bigtable_credentials import BigtableCredentialsConfig |
20 | 21 | from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset |
21 | 22 | from google.adk.tools.bigtable.settings import BigtableToolSettings |
| 23 | +from google.adk.tools.google_tool import GoogleTool |
22 | 24 | import google.auth |
| 25 | +from google.cloud.bigtable.data.execute_query.metadata import SqlType |
23 | 26 |
|
24 | | -# Define an appropriate credential type |
25 | | -CREDENTIALS_TYPE = AuthCredentialTypes.OAUTH2 |
26 | | - |
| 27 | +# Define an appropriate credential type. |
| 28 | +# None for Application Default Credentials |
| 29 | +CREDENTIALS_TYPE = None |
27 | 30 |
|
28 | 31 | # Define Bigtable tool config with read capability set to allowed. |
29 | 32 | tool_settings = BigtableToolSettings() |
|
59 | 62 | credentials_config=credentials_config, bigtable_tool_settings=tool_settings |
60 | 63 | ) |
61 | 64 |
|
| 65 | +_BIGTABLE_PROJECT_ID = "google.com:cloud-bigtable-dev" |
| 66 | +_BIGTABLE_INSTANCE_ID = "annenguyen-bus-instance" |
| 67 | + |
| 68 | + |
| 69 | +def search_hotels_by_location( |
| 70 | + location_name: str, |
| 71 | + credentials: google.auth.credentials.Credentials, |
| 72 | + settings: BigtableToolSettings, |
| 73 | + tool_context: google.adk.tools.tool_context.ToolContext, |
| 74 | +): |
| 75 | + """Search hotels by location name. |
| 76 | +
|
| 77 | + This function takes a location name and returns a list of hotels |
| 78 | + in that area. |
| 79 | +
|
| 80 | + Args: |
| 81 | + location_name (str): The geographical location (e.g., city or town) for the |
| 82 | + hotel search. |
| 83 | + Example: { "location_name": "Basel" } |
| 84 | +
|
| 85 | + Returns: |
| 86 | + The hotels name, price tier. |
| 87 | + """ |
| 88 | + |
| 89 | + sql_template = """ |
| 90 | + SELECT |
| 91 | + TO_INT64(cf['id']) as id, |
| 92 | + CAST(cf['name'] AS STRING) AS name, |
| 93 | + CAST(cf['location'] AS STRING) AS location, |
| 94 | + CAST(cf['price_tier'] AS STRING) AS price_tier, |
| 95 | + CAST(cf['checkin_date'] AS STRING) AS checkin_date, |
| 96 | + CAST(cf['checkout_date'] AS STRING) AS checkout_date |
| 97 | + FROM hotels |
| 98 | + WHERE LOWER(CAST(cf['location'] AS STRING)) LIKE LOWER(CONCAT('%', @location_name, '%')) |
| 99 | + """ |
| 100 | + return bigtable_query_tool.execute_sql( |
| 101 | + project_id=_BIGTABLE_PROJECT_ID, |
| 102 | + instance_id=_BIGTABLE_INSTANCE_ID, |
| 103 | + query=sql_template, |
| 104 | + credentials=credentials, |
| 105 | + settings=settings, |
| 106 | + tool_context=tool_context, |
| 107 | + parameters={"location": location_name}, |
| 108 | + parameter_types={"location": SqlType.String()}, |
| 109 | + ) |
| 110 | + |
| 111 | + |
62 | 112 | # The variable name `root_agent` determines what your root agent is for the |
63 | 113 | # debug CLI |
64 | 114 | root_agent = LlmAgent( |
|
72 | 122 | You are a data agent with access to several Bigtable tools. |
73 | 123 | Make use of those tools to answer the user's questions. |
74 | 124 | """, |
75 | | - tools=[bigtable_toolset], |
| 125 | + tools=[ |
| 126 | + bigtable_toolset, |
| 127 | + # Or, uncomment to use customized Bigtable tools. |
| 128 | + # GoogleTool( |
| 129 | + # func=search_hotels_by_location, |
| 130 | + # credentials_config=credentials_config, |
| 131 | + # tool_settings=tool_settings, |
| 132 | + # ), |
| 133 | + ], |
76 | 134 | ) |
0 commit comments