You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Compare commits
21 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| c6cf11cb74 | |||
| 2724819622 | |||
| 2f4f5611bd | |||
| 4df926388b | |||
| d4dc645478 | |||
| 7d5c6b9acf | |||
| aab2504ebd | |||
| 391628fcdc | |||
| 409df1378f | |||
| 2a901d12f4 | |||
| c850da3a07 | |||
| ed37e343f0 | |||
| ce3418a69d | |||
| fe1fc75c15 | |||
| dc4975dea9 | |||
| aeaec859bf | |||
| 4a842c5a13 | |||
| 0bdba30263 | |||
| 6a5eac0bdc | |||
| 0b73a6937b | |||
| ee39a89110 |
@@ -0,0 +1,18 @@
|
||||
# OAuth Sample
|
||||
|
||||
## Introduction
|
||||
|
||||
This sample data science agent uses Agent Engine Code Execution Sandbox to execute LLM generated code.
|
||||
|
||||
|
||||
## How to use
|
||||
|
||||
* 1. Follow https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/code-execution/overview to create a code execution sandbox environment.
|
||||
|
||||
* 2. Replace the SANDBOX_RESOURCE_NAME with the one you just created. If you dont want to create a new sandbox environment directly, the Agent Engine Code Execution Sandbox will create one for you by default using the AGENT_ENGINE_RESOURCE_NAME you specified, however, please ensure to clean up sandboxes after use, otherwise, it will consume quotas.
|
||||
|
||||
|
||||
## Sample prompt
|
||||
|
||||
* Can you write a function that calculates the sum from 1 to 100.
|
||||
* The dataset is given as below. Store,Date,Weekly_Sales,Holiday_Flag,Temperature,Fuel_Price,CPI,Unemployment Store 1,2023-06-01,1000,0,70,3.0,200,5 Store 2,2023-06-02,1200,1,80,3.5,210,6 Store 3,2023-06-03,1400,0,90,4.0,220,7 Store 4,2023-06-04,1600,1,70,4.5,230,8 Store 5,2023-06-05,1800,0,80,5.0,240,9 Store 6,2023-06-06,2000,1,90,5.5,250,10 Store 7,2023-06-07,2200,0,90,6.0,260,11 Plot a scatter plot showcasing the relationship between Weekly Sales and Temperature for each store, distinguishing stores with a Holiday Flag.
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# 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 . import agent
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""Data science agent."""
|
||||
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.code_executors.agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor
|
||||
|
||||
|
||||
def base_system_instruction():
|
||||
"""Returns: data science agent system instruction."""
|
||||
|
||||
return """
|
||||
# Guidelines
|
||||
|
||||
**Objective:** Assist the user in achieving their data analysis goals within the context of a Python Colab notebook, **with emphasis on avoiding assumptions and ensuring accuracy.** Reaching that goal can involve multiple steps. When you need to generate code, you **don't** need to solve the goal in one go. Only generate the next step at a time.
|
||||
|
||||
**Code Execution:** All code snippets provided will be executed within the Colab environment.
|
||||
|
||||
**Statefulness:** All code snippets are executed and the variables stays in the environment. You NEVER need to re-initialize variables. You NEVER need to reload files. You NEVER need to re-import libraries.
|
||||
|
||||
**Output Visibility:** Always print the output of code execution to visualize results, especially for data exploration and analysis. For example:
|
||||
- To look a the shape of a pandas.DataFrame do:
|
||||
```tool_code
|
||||
print(df.shape)
|
||||
```
|
||||
The output will be presented to you as:
|
||||
```tool_outputs
|
||||
(49, 7)
|
||||
|
||||
```
|
||||
- To display the result of a numerical computation:
|
||||
```tool_code
|
||||
x = 10 ** 9 - 12 ** 5
|
||||
print(f'{{x=}}')
|
||||
```
|
||||
The output will be presented to you as:
|
||||
```tool_outputs
|
||||
x=999751168
|
||||
|
||||
```
|
||||
- You **never** generate ```tool_outputs yourself.
|
||||
- You can then use this output to decide on next steps.
|
||||
- Print just variables (e.g., `print(f'{{variable=}}')`.
|
||||
|
||||
**No Assumptions:** **Crucially, avoid making assumptions about the nature of the data or column names.** Base findings solely on the data itself. Always use the information obtained from `explore_df` to guide your analysis.
|
||||
|
||||
**Available files:** Only use the files that are available as specified in the list of available files.
|
||||
|
||||
**Data in prompt:** Some queries contain the input data directly in the prompt. You have to parse that data into a pandas DataFrame. ALWAYS parse all the data. NEVER edit the data that are given to you.
|
||||
|
||||
**Answerability:** Some queries may not be answerable with the available data. In those cases, inform the user why you cannot process their query and suggest what type of data would be needed to fulfill their request.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.0-flash-001",
|
||||
name="agent_engine_code_execution_agent",
|
||||
instruction=base_system_instruction() + """
|
||||
|
||||
|
||||
You need to assist the user with their queries by looking at the data and the context in the conversation.
|
||||
You final answer should summarize the code and code execution relevant to the user query.
|
||||
|
||||
You should include all pieces of data to answer the user query, such as the table from code execution results.
|
||||
If you cannot answer the question directly, you should follow the guidelines above to generate the next step.
|
||||
If the question can be answered directly with writing any code, you should do that.
|
||||
If you doesn't have enough data to answer the question, you should ask for clarification from the user.
|
||||
|
||||
You should NEVER install any package on your own like `pip install ...`.
|
||||
When plotting trends, you should make sure to sort and order the data by the x-axis.
|
||||
|
||||
|
||||
""",
|
||||
code_executor=AgentEngineSandboxCodeExecutor(
|
||||
# Replace with your sandbox resource name if you already have one.
|
||||
sandbox_resource_name="SANDBOX_RESOURCE_NAME",
|
||||
# "projects/vertex-agent-loadtest/locations/us-central1/reasoningEngines/6842889780301135872/sandboxEnvironments/6545148628569161728",
|
||||
# Replace with agent engine resource name used for creating sandbox if
|
||||
# sandbox_resource_name is not set.
|
||||
agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME",
|
||||
),
|
||||
)
|
||||
@@ -17,17 +17,19 @@ This agent aims to test the Langchain tool with Langchain's StructuredTool
|
||||
"""
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.tools.langchain_tool import LangchainTool
|
||||
from langchain.tools import tool
|
||||
from langchain_core.tools import tool
|
||||
from langchain_core.tools.structured import StructuredTool
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
async def add(x, y) -> int:
|
||||
"""Adds two numbers."""
|
||||
return x + y
|
||||
|
||||
|
||||
@tool
|
||||
def minus(x, y) -> int:
|
||||
"""Subtracts two numbers."""
|
||||
return x - y
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ import random
|
||||
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.examples.example import Example
|
||||
from google.adk.models.google_llm import Gemini
|
||||
from google.adk.tools.example_tool import ExampleTool
|
||||
from google.genai import types
|
||||
|
||||
@@ -28,6 +29,17 @@ def roll_die(sides: int) -> int:
|
||||
|
||||
roll_agent = Agent(
|
||||
name="roll_agent",
|
||||
model=Gemini(
|
||||
# model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
|
||||
model="gemini-live-2.5-flash-preview", # for AI studio key
|
||||
speech_config=types.SpeechConfig(
|
||||
voice_config=types.VoiceConfig(
|
||||
prebuilt_voice_config=types.PrebuiltVoiceConfig(
|
||||
voice_name="Kore",
|
||||
)
|
||||
)
|
||||
),
|
||||
),
|
||||
description="Handles rolling dice of different sizes.",
|
||||
instruction="""
|
||||
You are responsible for rolling dice based on the user's request.
|
||||
@@ -69,6 +81,17 @@ def check_prime(nums: list[int]) -> str:
|
||||
|
||||
prime_agent = Agent(
|
||||
name="prime_agent",
|
||||
model=Gemini(
|
||||
# model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
|
||||
model="gemini-live-2.5-flash-preview", # for AI studio key
|
||||
speech_config=types.SpeechConfig(
|
||||
voice_config=types.VoiceConfig(
|
||||
prebuilt_voice_config=types.PrebuiltVoiceConfig(
|
||||
voice_name="Puck",
|
||||
)
|
||||
)
|
||||
),
|
||||
),
|
||||
description="Handles checking if numbers are prime.",
|
||||
instruction="""
|
||||
You are responsible for checking whether numbers are prime.
|
||||
@@ -100,8 +123,17 @@ def get_current_weather(location: str):
|
||||
|
||||
root_agent = Agent(
|
||||
# find supported models here: https://google.github.io/adk-docs/get-started/streaming/quickstart-streaming/
|
||||
model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
|
||||
# model="gemini-live-2.5-flash-preview", # for AI studio key
|
||||
model=Gemini(
|
||||
# model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
|
||||
model="gemini-live-2.5-flash-preview", # for AI studio key
|
||||
speech_config=types.SpeechConfig(
|
||||
voice_config=types.VoiceConfig(
|
||||
prebuilt_voice_config=types.PrebuiltVoiceConfig(
|
||||
voice_name="Zephyr",
|
||||
)
|
||||
)
|
||||
),
|
||||
),
|
||||
name="root_agent",
|
||||
instruction="""
|
||||
You are a helpful assistant that can check time, roll dice and check if numbers are prime.
|
||||
|
||||
@@ -4,37 +4,45 @@
|
||||
|
||||
This sample tests and demos the OAuth support in ADK via two tools:
|
||||
|
||||
* 1. list_calendar_events
|
||||
* 1. list_calendar_events
|
||||
|
||||
This is a customized tool that calls Google Calendar API to list calendar events.
|
||||
It pass in the client id and client secrete to ADK and then get back the access token from ADK.
|
||||
And then it uses the access token to call calendar api.
|
||||
This is a customized tool that calls Google Calendar API to list calendar
|
||||
events. It pass in the client id and client secrete to ADK and then get back
|
||||
the access token from ADK. And then it uses the access token to call
|
||||
calendar api.
|
||||
|
||||
* 2. get_calendar_events
|
||||
* 2. get_calendar_events
|
||||
|
||||
This is an google calendar tool that calls Google Calendar API to get the details of a specific calendar.
|
||||
This tool is from the ADK built-in Google Calendar ToolSet.
|
||||
Everything is wrapped and the tool user just needs to pass in the client id and client secret.
|
||||
This is an google calendar tool that calls Google Calendar API to get the
|
||||
details of a specific calendar. This tool is from the ADK built-in Google
|
||||
Calendar ToolSet. Everything is wrapped and the tool user just needs to pass
|
||||
in the client id and client secret.
|
||||
|
||||
## How to use
|
||||
|
||||
* 1. Follow https://developers.google.com/identity/protocols/oauth2#1.-obtain-oauth-2.0-credentials-from-the-dynamic_data.setvar.console_name. to get your client id and client secret.
|
||||
Be sure to choose "web" as your client type.
|
||||
* 1. Follow
|
||||
https://developers.google.com/identity/protocols/oauth2#1.-obtain-oauth-2.0-credentials-from-the-dynamic_data.setvar.console_name.
|
||||
to get your client id and client secret. Be sure to choose "web" as your
|
||||
client type.
|
||||
|
||||
* 2. Configure your `.env` file to add two variables:
|
||||
* 2. Configure your `.env` file to add two variables:
|
||||
|
||||
* OAUTH_CLIENT_ID={your client id}
|
||||
* OAUTH_CLIENT_SECRET={your client secret}
|
||||
* OAUTH_CLIENT_ID={your client id}
|
||||
* OAUTH_CLIENT_SECRET={your client secret}
|
||||
|
||||
Note: don't create a separate `.env` file , instead put it to the same `.env` file that stores your Vertex AI or Dev ML credentials
|
||||
Note: don't create a separate `.env` file , instead put it to the same
|
||||
`.env` file that stores your Vertex AI or Dev ML credentials
|
||||
|
||||
* 3. Follow https://developers.google.com/identity/protocols/oauth2/web-server#creatingcred to add http://localhost/dev-ui/ to "Authorized redirect URIs".
|
||||
* 3. Follow
|
||||
https://developers.google.com/identity/protocols/oauth2/web-server#creatingcred
|
||||
to add http://localhost/dev-ui/ to "Authorized redirect URIs".
|
||||
|
||||
Note: localhost here is just a hostname that you use to access the dev ui, replace it with the actual hostname you use to access the dev ui.
|
||||
Note: localhost here is just a hostname that you use to access the dev ui,
|
||||
replace it with the actual hostname you use to access the dev ui.
|
||||
|
||||
* 4. For 1st run, allow popup for localhost in Chrome.
|
||||
* 4. For 1st run, allow popup for localhost in Chrome.
|
||||
|
||||
## Sample prompt
|
||||
|
||||
* `List all my today's meeting from 7am to 7pm.`
|
||||
* `Get the details of the first event.`
|
||||
* `List all my today's meeting from 7am to 7pm.`
|
||||
* `Get the details of the first event.`
|
||||
|
||||
@@ -1,22 +1,32 @@
|
||||
# Output Schema with Tools Sample Agent
|
||||
|
||||
This sample demonstrates how to use structured output (`output_schema`) alongside other tools in an ADK agent. Previously, this combination was not allowed, but now it's supported through a special processor that handles the interaction.
|
||||
This sample demonstrates how to use structured output (`output_schema`)
|
||||
alongside other tools in an ADK agent. Previously, this combination was not
|
||||
allowed, but now it's supported through a special processor that handles the
|
||||
interaction.
|
||||
|
||||
## How it Works
|
||||
|
||||
The agent combines:
|
||||
- **Tools**: `search_wikipedia` and `get_current_year` for gathering information
|
||||
- **Structured Output**: `PersonInfo` schema to ensure consistent response format
|
||||
|
||||
- **Tools**: `search_wikipedia` and `get_current_year` for gathering
|
||||
information
|
||||
- **Structured Output**: `PersonInfo` schema to ensure consistent response
|
||||
format
|
||||
|
||||
When both `output_schema` and `tools` are specified:
|
||||
1. ADK automatically adds a special `set_model_response` tool
|
||||
2. The model can use the regular tools for information gathering
|
||||
3. For the final response, the model uses `set_model_response` with structured data
|
||||
4. ADK extracts and validates the structured response
|
||||
|
||||
1. ADK automatically adds a special `set_model_response` tool
|
||||
2. The model can use the regular tools for information gathering
|
||||
3. For the final response, the model uses `set_model_response` with structured
|
||||
data
|
||||
4. ADK extracts and validates the structured response
|
||||
|
||||
## Expected Response Format
|
||||
|
||||
The agent will return information in this structured format for user query "Tell me about Albert Einstein":
|
||||
The agent will return information in this structured format for user query
|
||||
|
||||
> Tell me about Albert Einstein.
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -30,7 +40,7 @@ The agent will return information in this structured format for user query "Tell
|
||||
|
||||
## Key Features Demonstrated
|
||||
|
||||
1. **Tool Usage**: Agent can search Wikipedia and get current year
|
||||
2. **Structured Output**: Response follows strict PersonInfo schema
|
||||
3. **Validation**: ADK validates the response matches the schema
|
||||
4. **Flexibility**: Works with any combination of tools and output schemas
|
||||
1. **Tool Usage**: Agent can search Wikipedia and get current year
|
||||
2. **Structured Output**: Response follows strict PersonInfo schema
|
||||
3. **Validation**: ADK validates the response matches the schema
|
||||
4. **Flexibility**: Works with any combination of tools and output schemas
|
||||
|
||||
+3
-3
@@ -32,7 +32,7 @@ dependencies = [
|
||||
"click>=8.1.8, <9.0.0", # For CLI tools
|
||||
"fastapi>=0.115.0, <1.0.0", # FastAPI framework
|
||||
"google-api-python-client>=2.157.0, <3.0.0", # Google API client discovery
|
||||
"google-cloud-aiplatform[agent_engines]>=1.112.0, <2.0.0",# For VertexAI integrations, e.g. example store.
|
||||
"google-cloud-aiplatform[agent_engines]>=1.121.0, <2.0.0",# For VertexAI integrations, e.g. example store.
|
||||
"google-cloud-bigtable>=2.32.0", # For Bigtable database
|
||||
"google-cloud-discoveryengine>=0.13.12, <0.14.0", # For Discovery Engine Search Tool
|
||||
"google-cloud-secret-manager>=2.22.0, <3.0.0", # Fetching secrets in RestAPI Tool
|
||||
@@ -114,7 +114,7 @@ test = [
|
||||
"anthropic>=0.43.0", # For anthropic model tests
|
||||
"kubernetes>=29.0.0", # For GkeCodeExecutor
|
||||
"langchain-community>=0.3.17",
|
||||
"langgraph>=0.2.60, <= 0.4.10", # For LangGraphAgent
|
||||
"langgraph>=0.2.60, <0.4.8", # For LangGraphAgent
|
||||
"litellm>=1.75.5, <2.0.0", # For LiteLLM tests
|
||||
"llama-index-readers-file>=0.4.0", # For retrieval tests
|
||||
"openai>=1.100.2", # For LiteLLM
|
||||
@@ -144,7 +144,7 @@ extensions = [
|
||||
"crewai[tools];python_version>='3.10'", # For CrewaiTool
|
||||
"docker>=7.0.0", # For ContainerCodeExecutor
|
||||
"kubernetes>=29.0.0", # For GkeCodeExecutor
|
||||
"langgraph>=0.2.60", # For LangGraphAgent
|
||||
"langgraph>=0.2.60, <0.4.8", # For LangGraphAgent
|
||||
"litellm>=1.75.5", # For LiteLlm class. Currently has OpenAI limitations. TODO: once LiteLlm fix it
|
||||
"llama-index-readers-file>=0.4.0", # For retrieval using LlamaIndex.
|
||||
"llama-index-embeddings-google-genai>=0.3.0",# For files retrieval using LlamaIndex.
|
||||
|
||||
@@ -282,27 +282,22 @@ class BaseAgent(BaseModel):
|
||||
Event: the events generated by the agent.
|
||||
"""
|
||||
|
||||
async def _run_with_trace() -> AsyncGenerator[Event, None]:
|
||||
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
|
||||
ctx = self._create_invocation_context(parent_context)
|
||||
tracing.trace_agent_invocation(span, self, ctx)
|
||||
if event := await self._handle_before_agent_callback(ctx):
|
||||
yield event
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
|
||||
ctx = self._create_invocation_context(parent_context)
|
||||
tracing.trace_agent_invocation(span, self, ctx)
|
||||
if event := await self._handle_before_agent_callback(ctx):
|
||||
yield event
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
|
||||
async with Aclosing(self._run_async_impl(ctx)) as agen:
|
||||
async for event in agen:
|
||||
yield event
|
||||
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
|
||||
if event := await self._handle_after_agent_callback(ctx):
|
||||
async with Aclosing(self._run_async_impl(ctx)) as agen:
|
||||
async for event in agen:
|
||||
yield event
|
||||
|
||||
async with Aclosing(_run_with_trace()) as agen:
|
||||
async for event in agen:
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
|
||||
if event := await self._handle_after_agent_callback(ctx):
|
||||
yield event
|
||||
|
||||
@final
|
||||
@@ -320,24 +315,19 @@ class BaseAgent(BaseModel):
|
||||
Event: the events generated by the agent.
|
||||
"""
|
||||
|
||||
async def _run_with_trace() -> AsyncGenerator[Event, None]:
|
||||
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
|
||||
ctx = self._create_invocation_context(parent_context)
|
||||
tracing.trace_agent_invocation(span, self, ctx)
|
||||
if event := await self._handle_before_agent_callback(ctx):
|
||||
yield event
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
|
||||
ctx = self._create_invocation_context(parent_context)
|
||||
tracing.trace_agent_invocation(span, self, ctx)
|
||||
if event := await self._handle_before_agent_callback(ctx):
|
||||
yield event
|
||||
if ctx.end_invocation:
|
||||
return
|
||||
|
||||
async with Aclosing(self._run_live_impl(ctx)) as agen:
|
||||
async for event in agen:
|
||||
yield event
|
||||
|
||||
if event := await self._handle_after_agent_callback(ctx):
|
||||
async with Aclosing(self._run_live_impl(ctx)) as agen:
|
||||
async for event in agen:
|
||||
yield event
|
||||
|
||||
async with Aclosing(_run_with_trace()) as agen:
|
||||
async for event in agen:
|
||||
if event := await self._handle_after_agent_callback(ctx):
|
||||
yield event
|
||||
|
||||
async def _run_async_impl(
|
||||
|
||||
@@ -143,7 +143,7 @@ Examples:
|
||||
|
||||
```
|
||||
# tools.py
|
||||
my_mcp_toolset = MCPToolset(
|
||||
my_mcp_toolset = McpToolset(
|
||||
connection_params=StdioServerParameters(
|
||||
command="npx",
|
||||
args=["-y", "@notionhq/notion-mcp-server"],
|
||||
|
||||
@@ -35,7 +35,10 @@ class StreamingMode(Enum):
|
||||
|
||||
|
||||
class RunConfig(BaseModel):
|
||||
"""Configs for runtime behavior of agents."""
|
||||
"""Configs for runtime behavior of agents.
|
||||
|
||||
The configs here will be overriden by agent-specific configurations.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
extra='forbid',
|
||||
|
||||
@@ -63,6 +63,7 @@ from ..agents.run_config import StreamingMode
|
||||
from ..apps.app import App
|
||||
from ..artifacts.base_artifact_service import BaseArtifactService
|
||||
from ..auth.credential_service.base_credential_service import BaseCredentialService
|
||||
from ..errors.already_exists_error import AlreadyExistsError
|
||||
from ..errors.not_found_error import NotFoundError
|
||||
from ..evaluation.base_eval_service import InferenceConfig
|
||||
from ..evaluation.base_eval_service import InferenceRequest
|
||||
@@ -486,6 +487,12 @@ class AdkWebServer:
|
||||
self.runner_dict[app_name] = runner
|
||||
return runner
|
||||
|
||||
def _get_root_agent(self, agent_or_app: BaseAgent | App) -> BaseAgent:
|
||||
"""Extract root agent from either a BaseAgent or App object."""
|
||||
if isinstance(agent_or_app, App):
|
||||
return agent_or_app.root_agent
|
||||
return agent_or_app
|
||||
|
||||
def _create_runner(self, agentic_app: App) -> Runner:
|
||||
"""Create a runner with common services."""
|
||||
return Runner(
|
||||
@@ -577,6 +584,33 @@ class AdkWebServer:
|
||||
"Failed to write runtime config file %s: %s", runtime_config_path, e
|
||||
)
|
||||
|
||||
async def _create_session(
|
||||
self,
|
||||
*,
|
||||
app_name: str,
|
||||
user_id: str,
|
||||
session_id: Optional[str] = None,
|
||||
state: Optional[dict[str, Any]] = None,
|
||||
) -> Session:
|
||||
try:
|
||||
session = await self.session_service.create_session(
|
||||
app_name=app_name,
|
||||
user_id=user_id,
|
||||
state=state,
|
||||
session_id=session_id,
|
||||
)
|
||||
logger.info("New session created: %s", session.id)
|
||||
return session
|
||||
except AlreadyExistsError as e:
|
||||
raise HTTPException(
|
||||
status_code=409, detail=f"Session already exists: {session_id}"
|
||||
) from e
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Internal server error during session creation: %s", e, exc_info=True
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||||
|
||||
def get_fast_api_app(
|
||||
self,
|
||||
lifespan: Optional[Lifespan[FastAPI]] = None,
|
||||
@@ -734,20 +768,12 @@ class AdkWebServer:
|
||||
session_id: str,
|
||||
state: Optional[dict[str, Any]] = None,
|
||||
) -> Session:
|
||||
if (
|
||||
await self.session_service.get_session(
|
||||
app_name=app_name, user_id=user_id, session_id=session_id
|
||||
)
|
||||
is not None
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=409, detail=f"Session already exists: {session_id}"
|
||||
)
|
||||
session = await self.session_service.create_session(
|
||||
app_name=app_name, user_id=user_id, state=state, session_id=session_id
|
||||
return await self._create_session(
|
||||
app_name=app_name,
|
||||
user_id=user_id,
|
||||
state=state,
|
||||
session_id=session_id,
|
||||
)
|
||||
logger.info("New session created: %s", session_id)
|
||||
return session
|
||||
|
||||
@app.post(
|
||||
"/apps/{app_name}/users/{user_id}/sessions",
|
||||
@@ -759,18 +785,9 @@ class AdkWebServer:
|
||||
req: Optional[CreateSessionRequest] = None,
|
||||
) -> Session:
|
||||
if not req:
|
||||
return await self.session_service.create_session(
|
||||
app_name=app_name, user_id=user_id
|
||||
)
|
||||
return await self._create_session(app_name=app_name, user_id=user_id)
|
||||
|
||||
if req.session_id and await self.session_service.get_session(
|
||||
app_name=app_name, user_id=user_id, session_id=req.session_id
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=409, detail=f"Session already exists: {req.session_id}"
|
||||
)
|
||||
|
||||
session = await self.session_service.create_session(
|
||||
session = await self._create_session(
|
||||
app_name=app_name,
|
||||
user_id=user_id,
|
||||
state=req.state,
|
||||
@@ -933,9 +950,8 @@ class AdkWebServer:
|
||||
|
||||
# Populate the session with initial session state.
|
||||
agent_or_app = self.agent_loader.load_agent(app_name)
|
||||
if isinstance(agent_or_app, App):
|
||||
agent_or_app = agent_or_app.root_agent
|
||||
initial_session_state = create_empty_state(agent_or_app)
|
||||
root_agent = self._get_root_agent(agent_or_app)
|
||||
initial_session_state = create_empty_state(root_agent)
|
||||
|
||||
new_eval_case = EvalCase(
|
||||
eval_id=req.eval_id,
|
||||
@@ -1096,7 +1112,8 @@ class AdkWebServer:
|
||||
status_code=400, detail=f"Eval set `{eval_set_id}` not found."
|
||||
)
|
||||
|
||||
root_agent = self.agent_loader.load_agent(app_name)
|
||||
agent_or_app = self.agent_loader.load_agent(app_name)
|
||||
root_agent = self._get_root_agent(agent_or_app)
|
||||
|
||||
eval_case_results = []
|
||||
|
||||
@@ -1437,13 +1454,7 @@ class AdkWebServer:
|
||||
function_calls = event.get_function_calls()
|
||||
function_responses = event.get_function_responses()
|
||||
agent_or_app = self.agent_loader.load_agent(app_name)
|
||||
# The loader may return an App; unwrap to its root agent so the graph builder
|
||||
# receives a BaseAgent instance.
|
||||
root_agent = (
|
||||
agent_or_app.root_agent
|
||||
if isinstance(agent_or_app, App)
|
||||
else agent_or_app
|
||||
)
|
||||
root_agent = self._get_root_agent(agent_or_app)
|
||||
dot_graph = None
|
||||
if function_calls:
|
||||
function_call_highlights = []
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -127,7 +127,7 @@ def conformance():
|
||||
pass
|
||||
|
||||
|
||||
@conformance.command("create", cls=HelpfulCommand)
|
||||
@conformance.command("record", cls=HelpfulCommand)
|
||||
@click.argument(
|
||||
"paths",
|
||||
nargs=-1,
|
||||
@@ -136,7 +136,7 @@ def conformance():
|
||||
),
|
||||
)
|
||||
@click.pass_context
|
||||
def cli_conformance_create(
|
||||
def cli_conformance_record(
|
||||
ctx,
|
||||
paths: tuple[str, ...],
|
||||
):
|
||||
@@ -162,7 +162,7 @@ def cli_conformance_create(
|
||||
"""
|
||||
|
||||
try:
|
||||
from .conformance.cli_create import run_conformance_create
|
||||
from .conformance.cli_record import run_conformance_record
|
||||
except ImportError as e:
|
||||
click.secho(
|
||||
f"Error: Missing conformance testing dependencies: {e}",
|
||||
@@ -178,7 +178,7 @@ def cli_conformance_create(
|
||||
|
||||
# Default to tests/ directory if no paths provided
|
||||
test_paths = [Path(p) for p in paths] if paths else [Path("tests").resolve()]
|
||||
asyncio.run(run_conformance_create(test_paths))
|
||||
asyncio.run(run_conformance_record(test_paths))
|
||||
|
||||
|
||||
@conformance.command("test", cls=HelpfulCommand)
|
||||
@@ -549,6 +549,7 @@ def cli_eval(
|
||||
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
|
||||
from ..evaluation.local_eval_sets_manager import load_eval_set_from_file
|
||||
from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
|
||||
from ..evaluation.user_simulator_provider import UserSimulatorProvider
|
||||
from .cli_eval import _collect_eval_results
|
||||
from .cli_eval import _collect_inferences
|
||||
from .cli_eval import get_root_agent
|
||||
@@ -638,11 +639,16 @@ def cli_eval(
|
||||
)
|
||||
)
|
||||
|
||||
user_simulator_provider = UserSimulatorProvider(
|
||||
user_simulator_config=eval_config.user_simulator_config
|
||||
)
|
||||
|
||||
try:
|
||||
eval_service = LocalEvalService(
|
||||
root_agent=root_agent,
|
||||
eval_sets_manager=eval_sets_manager,
|
||||
eval_set_results_manager=eval_set_results_manager,
|
||||
user_simulator_provider=user_simulator_provider,
|
||||
)
|
||||
|
||||
inference_results = asyncio.run(
|
||||
|
||||
+1
-1
@@ -105,7 +105,7 @@ async def _create_conformance_test_files(
|
||||
return generated_session_file
|
||||
|
||||
|
||||
async def run_conformance_create(paths: list[Path]) -> None:
|
||||
async def run_conformance_record(paths: list[Path]) -> None:
|
||||
"""Generate conformance tests from TestCaseInput files.
|
||||
|
||||
Args:
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user