|
| 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 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import inspect |
| 18 | +from typing import Any |
| 19 | +from typing import AsyncGenerator |
| 20 | +from typing import Awaitable |
| 21 | +from typing import Callable |
| 22 | +from typing import Dict |
| 23 | +from typing import final |
| 24 | +from typing import List |
| 25 | +from typing import Literal |
| 26 | +from typing import Mapping |
| 27 | +from typing import Optional |
| 28 | +from typing import Type |
| 29 | +from typing import TYPE_CHECKING |
| 30 | +from typing import TypeVar |
| 31 | +from typing import Union |
| 32 | + |
| 33 | +from google.genai import types |
| 34 | +from opentelemetry import trace |
| 35 | +from pydantic import BaseModel |
| 36 | +from pydantic import ConfigDict |
| 37 | +from pydantic import Field |
| 38 | +from pydantic import field_validator |
| 39 | +from pydantic import model_validator |
| 40 | +from typing_extensions import override |
| 41 | +from typing_extensions import TypeAlias |
| 42 | + |
| 43 | +from ..events.event import Event |
| 44 | +from ..utils.feature_decorator import working_in_progress |
| 45 | +from .callback_context import CallbackContext |
| 46 | +from .common_configs import CodeConfig |
| 47 | + |
| 48 | +if TYPE_CHECKING: |
| 49 | + from .invocation_context import InvocationContext |
| 50 | + |
| 51 | + |
| 52 | +TBaseAgentConfig = TypeVar('TBaseAgentConfig', bound='BaseAgentConfig') |
| 53 | + |
| 54 | + |
| 55 | +class SubAgentConfig(BaseModel): |
| 56 | + """The config for a sub-agent.""" |
| 57 | + |
| 58 | + model_config = ConfigDict(extra='forbid') |
| 59 | + |
| 60 | + config: Optional[str] = None |
| 61 | + """The YAML config file path of the sub-agent. |
| 62 | +
|
| 63 | + Only one of `config` or `code` can be set. |
| 64 | +
|
| 65 | + Example: |
| 66 | +
|
| 67 | + ``` |
| 68 | + sub_agents: |
| 69 | + - config: search_agent.yaml |
| 70 | + - config: my_library/my_custom_agent.yaml |
| 71 | + ``` |
| 72 | + """ |
| 73 | + |
| 74 | + code: Optional[str] = None |
| 75 | + """The agent instance defined in the code. |
| 76 | +
|
| 77 | + Only one of `config` or `code` can be set. |
| 78 | +
|
| 79 | + Example: |
| 80 | +
|
| 81 | + For the following agent defined in Python code: |
| 82 | +
|
| 83 | + ``` |
| 84 | + # my_library/custom_agents.py |
| 85 | + from google.adk.agents.llm_agent import LlmAgent |
| 86 | +
|
| 87 | + my_custom_agent = LlmAgent( |
| 88 | + name="my_custom_agent", |
| 89 | + instruction="You are a helpful custom agent.", |
| 90 | + model="gemini-2.0-flash", |
| 91 | + ) |
| 92 | + ``` |
| 93 | +
|
| 94 | + The yaml config should be: |
| 95 | +
|
| 96 | + ``` |
| 97 | + sub_agents: |
| 98 | + - code: my_library.custom_agents.my_custom_agent |
| 99 | + ``` |
| 100 | + """ |
| 101 | + |
| 102 | + @model_validator(mode='after') |
| 103 | + def validate_exactly_one_field(self): |
| 104 | + code_provided = self.code is not None |
| 105 | + config_provided = self.config is not None |
| 106 | + |
| 107 | + if code_provided and config_provided: |
| 108 | + raise ValueError('Only one of code or config should be provided') |
| 109 | + if not code_provided and not config_provided: |
| 110 | + raise ValueError('Exactly one of code or config must be provided') |
| 111 | + |
| 112 | + return self |
| 113 | + |
| 114 | + |
| 115 | +@working_in_progress('BaseAgentConfig is not ready for use.') |
| 116 | +class BaseAgentConfig(BaseModel): |
| 117 | + """The config for the YAML schema of a BaseAgent. |
| 118 | +
|
| 119 | + Do not use this class directly. It's the base class for all agent configs. |
| 120 | + """ |
| 121 | + |
| 122 | + model_config = ConfigDict( |
| 123 | + extra='allow', |
| 124 | + ) |
| 125 | + |
| 126 | + agent_class: Union[Literal['BaseAgent'], str] = 'BaseAgent' |
| 127 | + """Required. The class of the agent. The value is used to differentiate |
| 128 | + among different agent classes.""" |
| 129 | + |
| 130 | + name: str |
| 131 | + """Required. The name of the agent.""" |
| 132 | + |
| 133 | + description: str = '' |
| 134 | + """Optional. The description of the agent.""" |
| 135 | + |
| 136 | + sub_agents: Optional[List[SubAgentConfig]] = None |
| 137 | + """Optional. The sub-agents of the agent.""" |
| 138 | + |
| 139 | + before_agent_callbacks: Optional[List[CodeConfig]] = None |
| 140 | + """Optional. The before_agent_callbacks of the agent. |
| 141 | +
|
| 142 | + Example: |
| 143 | +
|
| 144 | + ``` |
| 145 | + before_agent_callbacks: |
| 146 | + - name: my_library.security_callbacks.before_agent_callback |
| 147 | + ``` |
| 148 | + """ |
| 149 | + |
| 150 | + after_agent_callbacks: Optional[List[CodeConfig]] = None |
| 151 | + """Optional. The after_agent_callbacks of the agent.""" |
| 152 | + |
| 153 | + def to_agent_config( |
| 154 | + self, custom_agent_config_cls: Type[TBaseAgentConfig] |
| 155 | + ) -> TBaseAgentConfig: |
| 156 | + """Converts this config to the concrete agent config type. |
| 157 | +
|
| 158 | + NOTE: this is for ADK framework use only. |
| 159 | + """ |
| 160 | + return custom_agent_config_cls.model_validate(self.model_dump()) |
0 commit comments