diff --git a/veadk/agent.py b/veadk/agent.py index 82e3eb1e..5defd0f3 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -56,6 +56,11 @@ from veadk.utils.patches import patch_asyncio, patch_tracer from veadk.utils.misc import check_litellm_version from veadk.version import VERSION +from veadk.prompts.prompt_example import ( + ExampleTool, + _convert_to_adk_examples, + AgentExample, +) patch_tracer() patch_asyncio() @@ -140,6 +145,8 @@ class Agent(LlmAgent): enable_authz: bool = False + examples: list[AgentExample] = Field(default_factory=list) + def model_post_init(self, __context: Any) -> None: super().model_post_init(None) # for sub_agents init @@ -258,6 +265,11 @@ def model_post_init(self, __context: Any) -> None: if self.prompt_manager: self.instruction = self.prompt_manager.get_prompt + if self.examples: + adk_examples = _convert_to_adk_examples(self.examples) + self.tools.append(ExampleTool(adk_examples)) + logger.info(f"Added {len(self.examples)} examples to agent") + logger.info(f"VeADK version: {VERSION}") logger.info(f"{self.__class__.__name__} `{self.name}` init done.") diff --git a/veadk/prompts/prompt_example.py b/veadk/prompts/prompt_example.py new file mode 100644 index 00000000..8943c983 --- /dev/null +++ b/veadk/prompts/prompt_example.py @@ -0,0 +1,57 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Union +from pydantic import BaseModel, Field +from google.genai import types +from google.adk.examples.example import Example + + +class FunctionCallExample(BaseModel): + """Represents a function call output in an example.""" + + func_name: str + func_args: dict[str, Any] = Field(default_factory=dict) + + +class AgentExample(BaseModel): + """A few-shot example for the agent. + + Attributes: + input: User input text. + output: Expected output - text string or function call. + """ + + input: str + output: Union[str, FunctionCallExample] + + +def _convert_to_adk_examples(examples: list[AgentExample]) -> list[Example]: + """Convert AgentExample list to ADK Example list.""" + result = [] + for ex in examples: + input_content = types.Content( + role="user", parts=[types.Part.from_text(text=ex.input)] + ) + if isinstance(ex.output, str): + output_parts = [types.Part.from_text(text=ex.output)] + else: + output_parts = [ + types.Part.from_function_call( + name=ex.output.func_name, args=ex.output.func_args + ) + ] + output_content = [types.Content(role="model", parts=output_parts)] + result.append(Example(input=input_content, output=output_content)) + return result