|
| 1 | +"""A minimal example demonstrating ConditionalAgent. |
| 2 | +
|
| 3 | +This sample shows how to route user requests to different sub-agents based on a |
| 4 | +predicate evaluated against the current `InvocationContext`. |
| 5 | +
|
| 6 | +• If the user asks to *roll* a die → `roll_agent` is triggered. |
| 7 | +• Otherwise, the request is delegated to `prime_agent` for prime-number checks. |
| 8 | +""" |
| 9 | +from __future__ import annotations |
| 10 | + |
| 11 | +import random |
| 12 | +from typing import cast |
| 13 | + |
| 14 | +from google.adk.agents.conditional_agent import ConditionalAgent |
| 15 | +from google.adk.agents.invocation_context import InvocationContext |
| 16 | +from google.adk.agents.llm_agent import LlmAgent |
| 17 | +from google.genai import types |
| 18 | + |
| 19 | +# ----------------------------------------------------------------------------- |
| 20 | +# Helper tool functions |
| 21 | +# ----------------------------------------------------------------------------- |
| 22 | + |
| 23 | +def roll_die(sides: int) -> int: |
| 24 | + """Roll a die with a given number of *sides* and return the result.""" |
| 25 | + return random.randint(1, sides) |
| 26 | + |
| 27 | + |
| 28 | +def check_prime(nums: list[int]) -> str: |
| 29 | + """Return a formatted string indicating which numbers are prime.""" |
| 30 | + primes = [] |
| 31 | + for n in nums: |
| 32 | + n = int(n) |
| 33 | + if n <= 1: |
| 34 | + continue |
| 35 | + for i in range(2, int(n ** 0.5) + 1): |
| 36 | + if n % i == 0: |
| 37 | + break |
| 38 | + else: |
| 39 | + primes.append(n) |
| 40 | + return "No prime numbers found." if not primes else ", ".join(map(str, primes)) + " are prime numbers." |
| 41 | + |
| 42 | +# ----------------------------------------------------------------------------- |
| 43 | +# Sub-agents definitions |
| 44 | +# ----------------------------------------------------------------------------- |
| 45 | + |
| 46 | +roll_agent = LlmAgent( |
| 47 | + name="roll_agent", |
| 48 | + description="Handles rolling dice of different sizes.", |
| 49 | + model="gemini-2.0-flash", |
| 50 | + instruction=( |
| 51 | + """ |
| 52 | + You are responsible for rolling dice based on the user's request.\n |
| 53 | + When asked to roll a die, call the `roll_die` tool with the number of |
| 54 | + sides. Do **not** decide the outcome yourself – always use the tool. |
| 55 | + """ |
| 56 | + ), |
| 57 | + tools=[roll_die], |
| 58 | + generate_content_config=types.GenerateContentConfig( |
| 59 | + safety_settings=[ |
| 60 | + types.SafetySetting( |
| 61 | + category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, |
| 62 | + threshold=types.HarmBlockThreshold.OFF, |
| 63 | + ), |
| 64 | + ] |
| 65 | + ), |
| 66 | +) |
| 67 | + |
| 68 | +prime_agent = LlmAgent( |
| 69 | + name="prime_agent", |
| 70 | + description="Checks whether provided numbers are prime.", |
| 71 | + model="gemini-2.0-flash", |
| 72 | + instruction=( |
| 73 | + """ |
| 74 | + You determine if numbers are prime.\n |
| 75 | + Whenever the user asks about prime numbers, call the `check_prime` tool |
| 76 | + with a list of integers and return its result. Never attempt to compute |
| 77 | + primes manually – always rely on the tool. |
| 78 | + """ |
| 79 | + ), |
| 80 | + tools=[check_prime], |
| 81 | + generate_content_config=types.GenerateContentConfig( |
| 82 | + safety_settings=[ |
| 83 | + types.SafetySetting( |
| 84 | + category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, |
| 85 | + threshold=types.HarmBlockThreshold.OFF, |
| 86 | + ), |
| 87 | + ] |
| 88 | + ), |
| 89 | +) |
| 90 | + |
| 91 | +# ----------------------------------------------------------------------------- |
| 92 | +# Predicate used by the ConditionalAgent |
| 93 | +# ----------------------------------------------------------------------------- |
| 94 | + |
| 95 | +def is_roll_request(ctx: InvocationContext) -> bool: |
| 96 | + """Return True if the last user message seems to be a *roll* request.""" |
| 97 | + if not ctx.user_content or not ctx.user_content.parts: |
| 98 | + return False |
| 99 | + text = cast(str, getattr(ctx.user_content.parts[0], "text", "")).lower() |
| 100 | + return "roll" in text |
| 101 | + |
| 102 | +# ----------------------------------------------------------------------------- |
| 103 | +# Root ConditionalAgent |
| 104 | +# ----------------------------------------------------------------------------- |
| 105 | + |
| 106 | +root_agent = ConditionalAgent( |
| 107 | + name="simple_conditional_agent", |
| 108 | + description="Routes to roll or prime agent based on user's intent.", |
| 109 | + sub_agents=[roll_agent, prime_agent], |
| 110 | + condition=is_roll_request, |
| 111 | +) |
0 commit comments