|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +"""Sales Data Assistant Agent demonstrating context offloading with artifacts. |
| 15 | +
|
| 16 | +This agent simulates querying large sales reports. To avoid cluttering |
| 17 | +the LLM context window with large amounts of data, queried reports are |
| 18 | +saved as artifacts rather than returned directly in function responses. |
| 19 | +Tools are used to inject artifact content into the LLM context only when |
| 20 | +needed: |
| 21 | +- QueryLargeDataTool injects content immediately after a report is generated. |
| 22 | +- CustomLoadArtifactsTool injects content when load_artifacts is called, and |
| 23 | + also provides artifact summaries to the LLM based on artifact metadata. |
| 24 | +""" |
| 25 | + |
| 26 | +import json |
| 27 | +import logging |
| 28 | +import random |
| 29 | + |
| 30 | +from google.adk import Agent |
| 31 | +from google.adk.apps import App |
| 32 | +from google.adk.models.llm_request import LlmRequest |
| 33 | +from google.adk.tools.function_tool import FunctionTool |
| 34 | +from google.adk.tools.load_artifacts_tool import LoadArtifactsTool |
| 35 | +from google.adk.tools.tool_context import ToolContext |
| 36 | +from google.genai import types |
| 37 | +from typing_extensions import override |
| 38 | + |
| 39 | +logger = logging.getLogger('google_adk.' + __name__) |
| 40 | + |
| 41 | + |
| 42 | +class CustomLoadArtifactsTool(LoadArtifactsTool): |
| 43 | + """A custom tool to load artifacts that also provides summaries. |
| 44 | +
|
| 45 | + This tool extends LoadArtifactsTool to read custom metadata from artifacts |
| 46 | + and provide summaries to the LLM in the system instructions, allowing the |
| 47 | + model to know what artifacts are available (e.g., "Sales report for APAC"). |
| 48 | + It also injects artifact content into the LLM request when load_artifacts |
| 49 | + is called by the model. |
| 50 | + """ |
| 51 | + |
| 52 | + @override |
| 53 | + async def _append_artifacts_to_llm_request( |
| 54 | + self, *, tool_context: ToolContext, llm_request: LlmRequest |
| 55 | + ): |
| 56 | + artifact_names = await tool_context.list_artifacts() |
| 57 | + if not artifact_names: |
| 58 | + return |
| 59 | + |
| 60 | + summaries = {} |
| 61 | + for name in artifact_names: |
| 62 | + version_info = await tool_context.get_artifact_version(name) |
| 63 | + if version_info and version_info.custom_metadata: |
| 64 | + summaries[name] = version_info.custom_metadata.get('summary') |
| 65 | + |
| 66 | + artifacts_with_summaries = [ |
| 67 | + f'{name}: {summaries.get(name)}' |
| 68 | + if name in summaries and summaries.get(name) |
| 69 | + else name |
| 70 | + for name in artifact_names |
| 71 | + ] |
| 72 | + |
| 73 | + # Tell the model about the available artifacts. |
| 74 | + llm_request.append_instructions([ |
| 75 | + f"""You have access to artifacts: {json.dumps(artifacts_with_summaries)}. |
| 76 | +If you need to answer a question that requires artifact content, first check if |
| 77 | +the content was very recently added to the conversation (e.g., in the last |
| 78 | +turn). If it is, use that content directly to answer. If the content is not |
| 79 | +available in the recent conversation history, you MUST call `load_artifacts` |
| 80 | +to retrieve it before answering. |
| 81 | +""" |
| 82 | + ]) |
| 83 | + |
| 84 | + # Attach the content of the artifacts if the model requests them. |
| 85 | + # This only adds the content to the model request, instead of the session. |
| 86 | + if llm_request.contents and llm_request.contents[-1].parts: |
| 87 | + function_response = llm_request.contents[-1].parts[0].function_response |
| 88 | + if function_response and function_response.name == 'load_artifacts': |
| 89 | + artifact_names = function_response.response['artifact_names'] |
| 90 | + if not artifact_names: |
| 91 | + return |
| 92 | + for artifact_name in artifact_names: |
| 93 | + # Try session-scoped first (default behavior) |
| 94 | + artifact = await tool_context.load_artifact(artifact_name) |
| 95 | + |
| 96 | + # If not found and name doesn't already have user: prefix, |
| 97 | + # try cross-session artifacts with user: prefix |
| 98 | + if artifact is None and not artifact_name.startswith('user:'): |
| 99 | + prefixed_name = f'user:{artifact_name}' |
| 100 | + artifact = await tool_context.load_artifact(prefixed_name) |
| 101 | + |
| 102 | + if artifact is None: |
| 103 | + logger.warning('Artifact "%s" not found, skipping', artifact_name) |
| 104 | + continue |
| 105 | + llm_request.contents.append( |
| 106 | + types.Content( |
| 107 | + role='user', |
| 108 | + parts=[ |
| 109 | + types.Part.from_text( |
| 110 | + text=f'Artifact {artifact_name} is:' |
| 111 | + ), |
| 112 | + artifact, |
| 113 | + ], |
| 114 | + ) |
| 115 | + ) |
| 116 | + |
| 117 | + |
| 118 | +async def query_large_data(query: str, tool_context: ToolContext) -> dict: |
| 119 | + """Generates a mock sales report for a given region and saves it as an artifact. |
| 120 | +
|
| 121 | + This function simulates querying a large dataset. It generates a mock report |
| 122 | + for North America, EMEA, or APAC, saves it as a text artifact, and includes |
| 123 | + a data summary in the artifact's custom metadata. |
| 124 | + Example queries: "Get sales data for North America", "EMEA sales report". |
| 125 | +
|
| 126 | + Args: |
| 127 | + query: The user query, expected to contain a region name. |
| 128 | + tool_context: The tool context for saving artifacts. |
| 129 | +
|
| 130 | + Returns: |
| 131 | + A dictionary containing a confirmation message and the artifact name. |
| 132 | + """ |
| 133 | + region = 'Unknown' |
| 134 | + if 'north america' in query.lower(): |
| 135 | + region = 'North America' |
| 136 | + elif 'emea' in query.lower(): |
| 137 | + region = 'EMEA' |
| 138 | + elif 'apac' in query.lower(): |
| 139 | + region = 'APAC' |
| 140 | + else: |
| 141 | + return { |
| 142 | + 'message': f"Sorry, I don't have data for query: {query}", |
| 143 | + 'artifact_name': None, |
| 144 | + } |
| 145 | + |
| 146 | + # simulate large data - Generate a mock sales report |
| 147 | + report_content = f"""SALES REPORT: {region} Q3 2025 |
| 148 | +========================================= |
| 149 | +Total Revenue: ${random.uniform(500, 2000):.2f}M |
| 150 | +Units Sold: {random.randint(100000, 500000)} |
| 151 | +Key Products: Gadget Pro, Widget Max, Thingy Plus |
| 152 | +Highlights: |
| 153 | +- Strong growth in Gadget Pro driven by new marketing campaign. |
| 154 | +- Widget Max sales are stable. |
| 155 | +- Thingy Plus saw a 15% increase in market share. |
| 156 | +
|
| 157 | +Regional Breakdown: |
| 158 | +""" + ''.join([ |
| 159 | + f'Sub-region {i+1} performance metric: {random.random()*100:.2f}\n' |
| 160 | + for i in range(500) |
| 161 | + ]) |
| 162 | + data_summary = f'Sales report for {region} Q3 2025' |
| 163 | + artifact_name = f"{region.replace(' ', '_')}_sales_report_q3_2025.txt" |
| 164 | + |
| 165 | + await tool_context.save_artifact( |
| 166 | + artifact_name, |
| 167 | + types.Part.from_text(text=report_content), |
| 168 | + custom_metadata={'summary': data_summary}, |
| 169 | + ) |
| 170 | + return { |
| 171 | + 'message': ( |
| 172 | + f'Sales data for {region} for Q3 2025 is saved as artifact' |
| 173 | + f" '{artifact_name}'." |
| 174 | + ), |
| 175 | + 'artifact_name': artifact_name, |
| 176 | + } |
| 177 | + |
| 178 | + |
| 179 | +class QueryLargeDataTool(FunctionTool): |
| 180 | + """A tool that queries large data and saves it as an artifact. |
| 181 | +
|
| 182 | + This tool wraps the query_large_data function. Its process_llm_request |
| 183 | + method checks if query_large_data was just called. If so, it loads the |
| 184 | + artifact that was just created and injects its content into the LLM |
| 185 | + request, so the model can use the data immediately in the next turn. |
| 186 | + """ |
| 187 | + |
| 188 | + def __init__(self): |
| 189 | + super().__init__(query_large_data) |
| 190 | + |
| 191 | + @override |
| 192 | + async def process_llm_request( |
| 193 | + self, |
| 194 | + *, |
| 195 | + tool_context: ToolContext, |
| 196 | + llm_request: LlmRequest, |
| 197 | + ) -> None: |
| 198 | + await super().process_llm_request( |
| 199 | + tool_context=tool_context, llm_request=llm_request |
| 200 | + ) |
| 201 | + if llm_request.contents and llm_request.contents[-1].parts: |
| 202 | + function_response = llm_request.contents[-1].parts[0].function_response |
| 203 | + if function_response and function_response.name == 'query_large_data': |
| 204 | + artifact_name = function_response.response.get('artifact_name') |
| 205 | + if artifact_name: |
| 206 | + artifact = await tool_context.load_artifact(artifact_name) |
| 207 | + if artifact: |
| 208 | + llm_request.contents.append( |
| 209 | + types.Content( |
| 210 | + role='user', |
| 211 | + parts=[ |
| 212 | + types.Part.from_text( |
| 213 | + text=f'Artifact {artifact_name} is:' |
| 214 | + ), |
| 215 | + artifact, |
| 216 | + ], |
| 217 | + ) |
| 218 | + ) |
| 219 | + |
| 220 | + |
| 221 | +root_agent = Agent( |
| 222 | + model='gemini-2.5-flash', |
| 223 | + name='context_offloading_with_artifact', |
| 224 | + description='An assistant for querying large sales reports.', |
| 225 | + instruction=""" |
| 226 | + You are a sales data assistant. You can query large sales reports by |
| 227 | + region (North America, EMEA, APAC) using the query_large_data tool. |
| 228 | + If you are asked to compare data between regions, make sure you have |
| 229 | + queried the data for all required regions first, and then use the |
| 230 | + load_artifacts tool if you need to access reports from previous turns. |
| 231 | + """, |
| 232 | + tools=[ |
| 233 | + QueryLargeDataTool(), |
| 234 | + CustomLoadArtifactsTool(), |
| 235 | + ], |
| 236 | + generate_content_config=types.GenerateContentConfig( |
| 237 | + safety_settings=[ |
| 238 | + types.SafetySetting( # avoid false alarm about rolling dice. |
| 239 | + category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, |
| 240 | + threshold=types.HarmBlockThreshold.OFF, |
| 241 | + ), |
| 242 | + ] |
| 243 | + ), |
| 244 | +) |
| 245 | + |
| 246 | + |
| 247 | +app = App( |
| 248 | + name='context_offloading_with_artifact', |
| 249 | + root_agent=root_agent, |
| 250 | +) |
0 commit comments