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
27 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2592f01eb6 | |||
| 6c0bf85042 | |||
| 7c8bc69dd0 | |||
| 5880109ab1 | |||
| 458d24e24e | |||
| 4a34501d38 | |||
| 43b484ff66 | |||
| 94d48fce32 | |||
| 2bd984adb3 | |||
| 59eda98eae | |||
| 8fb2be216f | |||
| 6b241f5ef2 | |||
| 86e7664006 | |||
| 0fe3870cd5 | |||
| fdc286a23c | |||
| 62fa4e513c | |||
| 4dd5434847 | |||
| d6c964e05e | |||
| 23d330eef1 | |||
| b8917bc80e | |||
| 3c51ee7f48 | |||
| f1ccc0cfca | |||
| 8e41f7f6c8 | |||
| da2c933b53 | |||
| 226e873b0f | |||
| 8afb99a078 | |||
| 688791396a |
@@ -1,5 +1,92 @@
|
||||
# Changelog
|
||||
|
||||
|
||||
## [1.22.1](https://github.com/google/adk-python/compare/v1.22.0...v1.22.1) (2026-01-09)
|
||||
|
||||
### Bug Fixes
|
||||
* Add back `adk migrate session` CLI ([8fb2be2](https://github.com/google/adk-python/commit/8fb2be216f11dabe7fa361a0402e5e6316878ad8)).
|
||||
* Escape database reserved keyword ([94d48fc](https://github.com/google/adk-python/commit/94d48fce32a1f07cef967d50e82f2b1975b4abd9)).
|
||||
|
||||
|
||||
## [1.22.0](https://github.com/google/adk-python/compare/v1.21.0...v1.22.0) (2026-01-08)
|
||||
|
||||
### Features
|
||||
|
||||
* **[Core]**
|
||||
* Make `LlmAgent.model` optional with a default fallback ([b287215](https://github.com/google/adk-python/commit/b28721508a41bf6bcfef52bbc042fb6193a32dfa)).
|
||||
* Support regex for allowed origins ([2ea6e51](https://github.com/google/adk-python/commit/2ea6e513cff61d3f330274725c66f82fce4ba259)).
|
||||
* Enable PROGRESSIVE_SSE_STREAMING feature by default ([0b1cff2](https://github.com/google/adk-python/commit/0b1cff2976d1c04acf3863f76107b05d1cec448f)).
|
||||
|
||||
* **[Evals]**
|
||||
* Add custom instructions support to LlmBackedUserSimulator ([a364388](https://github.com/google/adk-python/commit/a364388d9744969760fd87ed24d60793146c162a)).
|
||||
* Introduce a post-hoc, per-turn evaluator for user simulations ([e515e0f](https://github.com/google/adk-python/commit/e515e0f321a259016c5e5f6b388ecf02ae343ba7)).
|
||||
|
||||
* **[Tools]**
|
||||
* Expose mcps streamable http custom httpx factory parameter ([bfed19c](https://github.com/google/adk-python/commit/bfed19cd78298fc9f896da8ed82a756004e92094)).
|
||||
* Add a handwritten tool for Cloud Pub/Sub ([b6f6dcb](https://github.com/google/adk-python/commit/b6f6dcbeb465a775b9c38ace7a324ee2155d366f)).
|
||||
* Add `token_endpoint_auth_method` support to OAuth2 credentials ([8782a69](https://github.com/google/adk-python/commit/8782a695036aa0c1528027673868159143f925f0)).
|
||||
|
||||
* **[Services]**
|
||||
* Introduce new JSON-based database schema for DatabaseSessionService, which will be used for newly-created databases. A migration command and script are provided.([7e6ef71](https://github.com/google/adk-python/commit/7e6ef71eec8be2e804286cc4140d0cbdf84f1206) [ba91fea](https://github.com/google/adk-python/commit/ba91fea54136ab60f37c10b899c3648d0b0fa721) [ce64787](https://github.com/google/adk-python/commit/ce64787c3e1130d1678e408aa31011fc88590e15)).
|
||||
* Set log level when deploying to Agent Engine ([1f546df](https://github.com/google/adk-python/commit/1f546df35a1c18aeb3d2fc7a2ac66edf386027c5)).
|
||||
|
||||
* **[A2A]**
|
||||
* Update event_converter used in A2ARemote agent to use a2a_task.status.message only if parts are non-empty ([e4ee9d7](https://github.com/google/adk-python/commit/e4ee9d7c46b57eed8493539d8f539c042bdfae60)).
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* Add checks for event content and parts before accessing ([5912835](https://github.com/google/adk-python/commit/5912835c975673c8fc2fb315150f5ec29d685eac)).
|
||||
* Validate app name in `adk create` command ([742c926](https://github.com/google/adk-python/commit/742c9265a260a9c598a1f65e0996d926b4b9c022)).
|
||||
* Prevent .env files from overriding existing environment variables ([0827d12](https://github.com/google/adk-python/commit/0827d12ccd74feb24758f64f2884c9493001b4ca)).
|
||||
* Prevent ContextFilterPlugin from creating orphaned function responses ([e32f017](https://github.com/google/adk-python/commit/e32f017979e26a94b998311cafcde753fd29e44e)).
|
||||
* Update empty event check to include executable code and execution results ([688f48f](https://github.com/google/adk-python/commit/688f48fffb9df6ef18a692cd2ccaa7628f4c82a7)).
|
||||
* Make the BigQuery analytics plugin work with agents that don't have instructions such as the LoopAgent ([8bed01c](https://github.com/google/adk-python/commit/8bed01cbdc5961c0d219fd6389f492f1a4235de5)).
|
||||
* Label response as thought if task is immediately returned as working ([4f3b733](https://github.com/google/adk-python/commit/4f3b733074c867e68ca5d38720ccb1f3e0b0d960)).
|
||||
* Move and enhance the deprecation warning for the plugins argument in "_validate_runner_params" to the beginning of the function ([43270bc](https://github.com/google/adk-python/commit/43270bcb6197526ba5765f83d7e4fb88f213b8d3)).
|
||||
* Oauth refresh not triggered on token expiry ([69997cd](https://github.com/google/adk-python/commit/69997cd5ef44ee881a974bb36dc100e17ed6de2e)).
|
||||
* Fix double JSON encoding when saving eval set results ([fc4e3d6](https://github.com/google/adk-python/commit/fc4e3d6f607032259e68e91bcb1ad0815a03164e)).
|
||||
* Allow string values for ToolTrajectoryCriterion.match_type ([93d6e4c](https://github.com/google/adk-python/commit/93d6e4c888d5a2181e3c22da049d8be0d6ead70c)).
|
||||
* Fix inconsistent method signatures for evaluate_invocations ([0918b64](https://github.com/google/adk-python/commit/0918b647df6f88b95974d486a3161121a6514901)).
|
||||
* Honor the modalities parameter in adk api server for live API ([19de45b](https://github.com/google/adk-python/commit/19de45b3250d09b9ec16c45788e7d472b3e588c2)).
|
||||
* Filter out thought parts in lite_llm._get_content ([1ace8fc](https://github.com/google/adk-python/commit/1ace8fc6780bc25e2ef4222c73bc2558082b0a00)).
|
||||
* Rehydration of EventActions in StorageEvent.to_event ([838530e](https://github.com/google/adk-python/commit/838530ebe053e5193d4329c5a203ca3d096ff7be)).
|
||||
* Heal missing tool results before LiteLLM requests ([6b7386b](https://github.com/google/adk-python/commit/6b7386b7620bbc51cda8c1c6d9914549536640e6)).
|
||||
* Refine Ollama content flattening and provider checks ([c6f389d](https://github.com/google/adk-python/commit/c6f389d4bc4d2b91795003a3bd87ed1f1b854493)).
|
||||
* Add MIME type inference and default for file URIs in LiteLLM ([5c4bae7](https://github.com/google/adk-python/commit/5c4bae7ff2085c05b7f002f5fa368e9b48a752b1)).
|
||||
* Use mode='json' in model_dump to serialize bytes correctly when using telemetry ([96c5db5](https://github.com/google/adk-python/commit/96c5db5a07f7f851751ccd68f176dad1634885cb)).
|
||||
* Avoid local .adk storage in Cloud Run/GKE ([b30c2f4](https://github.com/google/adk-python/commit/b30c2f4e139e0d4410c5f8dd61acee2056ad06ea)).
|
||||
* Remove fallback to cached exchanged credential in _load_existing_credential ([1ae0e16](https://github.com/google/adk-python/commit/1ae0e16b2c1a3139b9c2b1c4a3e725833a6240be)).
|
||||
* Handle overriding of requirements when deploying to agent engine ([38a30a4](https://github.com/google/adk-python/commit/38a30a44d222fade8616f9d63410b1c2b6f84e1b)).
|
||||
* Built-in agents (names starting with "__") now use in-memory session storage instead of creating .adk folders in the agents directory ([e3bac1a](https://github.com/google/adk-python/commit/e3bac1ab8c724454fb433cc7e78416b61efe33ee)).
|
||||
* Change error_message column type to TEXT in DatabaseSessionService ([8335f35](https://github.com/google/adk-python/commit/8335f35015c7d4349bc4ac47dedbe99663b78e62)).
|
||||
* Add schema type sanitization to OpenAPI spec parser ([6dce7f8](https://github.com/google/adk-python/commit/6dce7f8a8f28de275b1119fc03219f1468bb883b)).
|
||||
* Prevent retry_on_errors from retrying asyncio.CancelledError ([30d3411](https://github.com/google/adk-python/commit/30d3411d603f12ca5bcdd2d71773d087f3191dba)).
|
||||
* Include back-ticks around the BQ asset names in the tools examples ([8789ad8](https://github.com/google/adk-python/commit/8789ad8f16dfa250fab607946250a2857a25d5ef)).
|
||||
* Fix issue with MCP tools throwing an error ([26e77e1](https://github.com/google/adk-python/commit/26e77e16947aed1abcfdd7f526532a708f1f073b)).
|
||||
* Exclude thought parts when merging agent output ([07bb164](https://github.com/google/adk-python/commit/07bb1647588a781e701257c4c379736537029ea0)).
|
||||
* Prepend "https://" to the MCP server url only if it doesn't already have a scheme ([71b3289](https://github.com/google/adk-python/commit/71b32890f5ab279e2bed1fd28c0f4693cba3f45e)).
|
||||
* Split SSE events with both content and artifactDelta in ADK Web Server ([084fcfa](https://github.com/google/adk-python/commit/084fcfaba52c4a6075397221dbe7aba2f2acd2d7)).
|
||||
* Propagate RunConfig custom metadata to all events ([e3db2d0](https://github.com/google/adk-python/commit/e3db2d0d8301748c63bad826f24692448dbd1c2c)).
|
||||
* Harden YAML builder tmp save/cleanup([6f259f0](https://github.com/google/adk-python/commit/6f259f08b3c45ad6050b8a93c9bd85913451ece6)).
|
||||
* Ignore adk-bot administrative actions in stale agent ([3ec7ae3](https://github.com/google/adk-python/commit/3ec7ae3b8d532ed4b60786201a78e980dfc56cf3)).
|
||||
* Only prepend "https://" to the MCP server url if it doesn't already have a scheme ([71b3289](https://github.com/google/adk-python/commit/71b32890f5ab279e2bed1fd28c0f4693cba3f45e)).
|
||||
* Check all content parts for emptiness in _contains_empty_content ([f35d129](https://github.com/google/adk-python/commit/f35d129b4c59d381e95418725d6eaa072ca7720a)).
|
||||
|
||||
### Improvements
|
||||
|
||||
* Remove unnecessary event loop creation in LiveRequstQueue constructor ([ecc9f18](https://github.com/google/adk-python/commit/ecc9f182e3bd25ee8eda8920d665e967517ca59a)).
|
||||
* Close database engines to avoid aiosqlite pytest hangs ([4ddb2cb](https://github.com/google/adk-python/commit/4ddb2cb2a8d1d026a43418b2dd698e6ea199594e)).
|
||||
* Add `override_feature_enabled` to override the default feature enable states ([a088506](https://github.com/google/adk-python/commit/a0885064b0cbef3b25484025da0748dc64320d4a)).
|
||||
* Move SQLite migration script to migration/ folder ([e8ab7da](https://github.com/google/adk-python/commit/e8ab7dafa96d5890a4fff919b9fa180993ef5830)).
|
||||
* Update latest Live Model names for sample agent ([f1eb1c0](https://github.com/google/adk-python/commit/f1eb1c0254802ef3aa64c76512e3104376291ec0)).
|
||||
* Update google-genai and google-cloud-aiplatform versions ([d58ea58](https://github.com/google/adk-python/commit/d58ea589ade822894f1482fd505a33d842755d9c)).
|
||||
* Introduce MetricInfoProvider interface, and refactor metric evaluators to use this interface to provide MetricInfo ([5b7c8c0](https://github.com/google/adk-python/commit/5b7c8c04d6e4a688c76fa517922488e3d96353a3)).
|
||||
* Update _flatten_ollama_content return type and add tests ([fcea86f](https://github.com/google/adk-python/commit/fcea86f58c95894bc9c1fb7ed12e36ddedaaa88a)).
|
||||
* Add disambiguation message to enterprise_search_tool ([8329fec](https://github.com/google/adk-python/commit/8329fec0fc6b6130ffd1f53a8a2e2ccc6e8f43ed)).
|
||||
* Add x-goog-user-project header to http calls in API Registry ([0088b0f](https://github.com/google/adk-python/commit/0088b0f3adb963dded692929c314d94709dcc211)).
|
||||
* Set the default response modality to AUDIO only ([a4b914b](https://github.com/google/adk-python/commit/a4b914b09fbab76834050a8c8f0eb335b12cfc34)).
|
||||
|
||||
|
||||
## [1.21.0](https://github.com/google/adk-python/compare/v1.20.0...v1.21.0) (2025-12-11)
|
||||
|
||||
### Features
|
||||
|
||||
@@ -114,7 +114,7 @@ async def _build_llm_agent_skills(agent: LlmAgent) -> List[AgentSkill]:
|
||||
id=agent.name,
|
||||
name='model',
|
||||
description=agent_description,
|
||||
examples=agent_examples,
|
||||
examples=_extract_inputs_from_examples(agent_examples),
|
||||
input_modes=_get_input_modes(agent),
|
||||
output_modes=_get_output_modes(agent),
|
||||
tags=['llm'],
|
||||
@@ -239,7 +239,7 @@ async def _build_non_llm_agent_skills(agent: BaseAgent) -> List[AgentSkill]:
|
||||
id=agent.name,
|
||||
name=agent_name,
|
||||
description=agent_description,
|
||||
examples=agent_examples,
|
||||
examples=_extract_inputs_from_examples(agent_examples),
|
||||
input_modes=_get_input_modes(agent),
|
||||
output_modes=_get_output_modes(agent),
|
||||
tags=[agent_type],
|
||||
@@ -350,6 +350,7 @@ def _build_llm_agent_description_with_instructions(agent: LlmAgent) -> str:
|
||||
|
||||
def _replace_pronouns(text: str) -> str:
|
||||
"""Replace pronouns and conjugate common verbs for agent description.
|
||||
|
||||
(e.g., "You are" -> "I am", "your" -> "my").
|
||||
"""
|
||||
pronoun_map = {
|
||||
@@ -460,6 +461,33 @@ def _get_default_description(agent: BaseAgent) -> str:
|
||||
return 'A custom agent'
|
||||
|
||||
|
||||
def _extract_inputs_from_examples(examples: Optional[list[dict]]) -> list[str]:
|
||||
"""Extracts only the input strings so they can be added to an AgentSkill."""
|
||||
if examples is None:
|
||||
return []
|
||||
|
||||
extracted_inputs = []
|
||||
for example in examples:
|
||||
example_input = example.get('input')
|
||||
if not example_input:
|
||||
continue
|
||||
|
||||
parts = example_input.get('parts')
|
||||
if parts is not None:
|
||||
part_texts = []
|
||||
for part in parts:
|
||||
text = part.get('text')
|
||||
if text is not None:
|
||||
part_texts.append(text)
|
||||
extracted_inputs.append('\n'.join(part_texts))
|
||||
else:
|
||||
text = example_input.get('text')
|
||||
if text is not None:
|
||||
extracted_inputs.append(text)
|
||||
|
||||
return extracted_inputs
|
||||
|
||||
|
||||
async def _extract_examples_from_agent(
|
||||
agent: BaseAgent,
|
||||
) -> Optional[List[Dict]]:
|
||||
|
||||
@@ -126,6 +126,7 @@ class RemoteA2aAgent(BaseAgent):
|
||||
a2a_request_meta_provider: Optional[
|
||||
Callable[[InvocationContext, A2AMessage], dict[str, Any]]
|
||||
] = None,
|
||||
full_history_when_stateless: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize RemoteA2aAgent.
|
||||
@@ -142,6 +143,10 @@ class RemoteA2aAgent(BaseAgent):
|
||||
a2a_request_meta_provider: Optional callable that takes InvocationContext
|
||||
and A2AMessage and returns a metadata object to attach to the A2A
|
||||
request.
|
||||
full_history_when_stateless: If True, stateless agents (those that do not
|
||||
return Tasks or context IDs) will receive all session events on every
|
||||
request. If False, the default behavior of sending only events since the
|
||||
last reply from the agent will be used.
|
||||
**kwargs: Additional arguments passed to BaseAgent
|
||||
|
||||
Raises:
|
||||
@@ -168,6 +173,7 @@ class RemoteA2aAgent(BaseAgent):
|
||||
self._a2a_part_converter = a2a_part_converter
|
||||
self._a2a_client_factory: Optional[A2AClientFactory] = a2a_client_factory
|
||||
self._a2a_request_meta_provider = a2a_request_meta_provider
|
||||
self._full_history_when_stateless = full_history_when_stateless
|
||||
|
||||
# Validate and store agent card reference
|
||||
if isinstance(agent_card, AgentCard):
|
||||
@@ -365,7 +371,14 @@ class RemoteA2aAgent(BaseAgent):
|
||||
if event.custom_metadata:
|
||||
metadata = event.custom_metadata
|
||||
context_id = metadata.get(A2A_METADATA_PREFIX + "context_id")
|
||||
break
|
||||
# Historical note: this behavior originally always applied, regardless
|
||||
# of whether the agent was stateful or stateless. However, only stateful
|
||||
# agents can be expected to have previous events in the remote session.
|
||||
# For backwards compatibility, we maintain this behavior when
|
||||
# _full_history_when_stateless is false (the default) or if the agent
|
||||
# is stateful (i.e. returned a context ID).
|
||||
if not self._full_history_when_stateless or context_id:
|
||||
break
|
||||
events_to_process.append(event)
|
||||
|
||||
for event in reversed(events_to_process):
|
||||
|
||||
@@ -61,6 +61,7 @@ class HttpAuth(BaseModelWithConfig):
|
||||
# Examples: 'basic', 'bearer'
|
||||
scheme: str
|
||||
credentials: HttpCredentials
|
||||
additional_headers: Optional[Dict[str, str]] = None
|
||||
|
||||
|
||||
class OAuth2Auth(BaseModelWithConfig):
|
||||
|
||||
@@ -64,7 +64,7 @@ COPY --chown=myuser:myuser "agents/{app_name}/" "/app/agents/{app_name}/"
|
||||
|
||||
EXPOSE {port}
|
||||
|
||||
CMD adk {command} --port={port} {host_option} {service_option} {trace_to_cloud_option} {allow_origins_option} {a2a_option} "/app/agents"
|
||||
CMD adk {command} --port={port} {host_option} {service_option} {trace_to_cloud_option} {otel_to_cloud_option} {allow_origins_option} {a2a_option} "/app/agents"
|
||||
"""
|
||||
|
||||
_AGENT_ENGINE_APP_TEMPLATE: Final[str] = """
|
||||
@@ -487,6 +487,7 @@ def to_cloud_run(
|
||||
temp_folder: str,
|
||||
port: int,
|
||||
trace_to_cloud: bool,
|
||||
otel_to_cloud: bool,
|
||||
with_ui: bool,
|
||||
log_level: str,
|
||||
verbosity: str,
|
||||
@@ -523,6 +524,8 @@ def to_cloud_run(
|
||||
temp_folder: The temp folder for the generated Cloud Run source files.
|
||||
port: The port of the ADK api server.
|
||||
trace_to_cloud: Whether to enable Cloud Trace.
|
||||
otel_to_cloud: Whether to enable exporting OpenTelemetry signals
|
||||
to Google Cloud.
|
||||
with_ui: Whether to deploy with UI.
|
||||
verbosity: The verbosity level of the CLI.
|
||||
adk_version: The ADK version to use in Cloud Run.
|
||||
@@ -580,6 +583,7 @@ def to_cloud_run(
|
||||
use_local_storage,
|
||||
),
|
||||
trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
|
||||
otel_to_cloud_option='--otel_to_cloud' if otel_to_cloud else '',
|
||||
allow_origins_option=allow_origins_option,
|
||||
adk_version=adk_version,
|
||||
host_option=host_option,
|
||||
@@ -956,6 +960,7 @@ def to_gke(
|
||||
temp_folder: str,
|
||||
port: int,
|
||||
trace_to_cloud: bool,
|
||||
otel_to_cloud: bool,
|
||||
with_ui: bool,
|
||||
log_level: str,
|
||||
adk_version: str,
|
||||
@@ -981,6 +986,8 @@ def to_gke(
|
||||
Dockerfile and deployment.yaml.
|
||||
port: The port of the ADK api server.
|
||||
trace_to_cloud: Whether to enable Cloud Trace.
|
||||
otel_to_cloud: Whether to enable exporting OpenTelemetry signals
|
||||
to Google Cloud.
|
||||
with_ui: Whether to deploy with UI.
|
||||
log_level: The logging level.
|
||||
adk_version: The ADK version to use in GKE.
|
||||
@@ -1051,6 +1058,7 @@ def to_gke(
|
||||
use_local_storage,
|
||||
),
|
||||
trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
|
||||
otel_to_cloud_option='--otel_to_cloud' if otel_to_cloud else '',
|
||||
allow_origins_option=allow_origins_option,
|
||||
adk_version=adk_version,
|
||||
host_option=host_option,
|
||||
|
||||
@@ -201,9 +201,11 @@ def pretty_print_eval_result(eval_result: EvalCaseResult):
|
||||
for r in metric_result.criterion.rubrics
|
||||
}
|
||||
for rubric_score in metric_result.details.rubric_scores:
|
||||
rubric = rubrics_by_id.get(rubric_score.rubric_id)
|
||||
rubric_text = rubrics_by_id.get(rubric_score.rubric_id)
|
||||
if not rubric_text:
|
||||
rubric_text = rubric_score.rubric_id
|
||||
click.echo(
|
||||
f"Rubric: {rubric}, "
|
||||
f"Rubric: {rubric_text}, "
|
||||
f"Score: {rubric_score.score}, "
|
||||
f"Reasoning: {rubric_score.rationale}"
|
||||
)
|
||||
@@ -243,6 +245,8 @@ def pretty_print_eval_result(eval_result: EvalCaseResult):
|
||||
}
|
||||
for rubric_score in metric_result.details.rubric_scores:
|
||||
rubric = rubrics_by_id.get(rubric_score.rubric_id)
|
||||
if not rubric:
|
||||
rubric = rubric_score.rubric_id
|
||||
row_data[f"Rubric: {rubric}"] = (
|
||||
f"Reasoning: {rubric_score.rationale}, "
|
||||
f"Score: {rubric_score.score}"
|
||||
|
||||
@@ -36,6 +36,7 @@ from . import cli_create
|
||||
from . import cli_deploy
|
||||
from .. import version
|
||||
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
|
||||
from ..sessions.migration import migration_runner
|
||||
from .cli import run_cli
|
||||
from .fast_api import get_fast_api_app
|
||||
from .utils import envs
|
||||
@@ -1370,6 +1371,13 @@ def cli_api_server(
|
||||
default=False,
|
||||
help="Optional. Whether to enable Cloud Trace for cloud run.",
|
||||
)
|
||||
@click.option(
|
||||
"--otel_to_cloud",
|
||||
is_flag=True,
|
||||
show_default=True,
|
||||
default=False,
|
||||
help="Optional. Whether to enable OpenTelemetry for Agent Engine.",
|
||||
)
|
||||
@click.option(
|
||||
"--with_ui",
|
||||
is_flag=True,
|
||||
@@ -1450,6 +1458,7 @@ def cli_deploy_cloud_run(
|
||||
temp_folder: str,
|
||||
port: int,
|
||||
trace_to_cloud: bool,
|
||||
otel_to_cloud: bool,
|
||||
with_ui: bool,
|
||||
adk_version: str,
|
||||
log_level: str,
|
||||
@@ -1528,6 +1537,7 @@ def cli_deploy_cloud_run(
|
||||
temp_folder=temp_folder,
|
||||
port=port,
|
||||
trace_to_cloud=trace_to_cloud,
|
||||
otel_to_cloud=otel_to_cloud,
|
||||
allow_origins=allow_origins,
|
||||
with_ui=with_ui,
|
||||
log_level=log_level,
|
||||
@@ -1544,6 +1554,47 @@ def cli_deploy_cloud_run(
|
||||
click.secho(f"Deploy failed: {e}", fg="red", err=True)
|
||||
|
||||
|
||||
@main.group()
|
||||
def migrate():
|
||||
"""ADK migration commands."""
|
||||
pass
|
||||
|
||||
|
||||
@migrate.command("session", cls=HelpfulCommand)
|
||||
@click.option(
|
||||
"--source_db_url",
|
||||
required=True,
|
||||
help=(
|
||||
"SQLAlchemy URL of source database in database session service, e.g."
|
||||
" sqlite:///source.db."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--dest_db_url",
|
||||
required=True,
|
||||
help=(
|
||||
"SQLAlchemy URL of destination database in database session service,"
|
||||
" e.g. sqlite:///dest.db."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--log_level",
|
||||
type=LOG_LEVELS,
|
||||
default="INFO",
|
||||
help="Optional. Set the logging level",
|
||||
)
|
||||
def cli_migrate_session(
|
||||
*, source_db_url: str, dest_db_url: str, log_level: str
|
||||
):
|
||||
"""Migrates a session database to the latest schema version."""
|
||||
logs.setup_adk_logger(getattr(logging, log_level.upper()))
|
||||
try:
|
||||
migration_runner.upgrade(source_db_url, dest_db_url)
|
||||
click.secho("Migration check and upgrade process finished.", fg="green")
|
||||
except Exception as e:
|
||||
click.secho(f"Migration failed: {e}", fg="red", err=True)
|
||||
|
||||
|
||||
@deploy.command("agent_engine")
|
||||
@click.option(
|
||||
"--api_key",
|
||||
@@ -1799,6 +1850,13 @@ def cli_deploy_agent_engine(
|
||||
default=False,
|
||||
help="Optional. Whether to enable Cloud Trace for GKE.",
|
||||
)
|
||||
@click.option(
|
||||
"--otel_to_cloud",
|
||||
is_flag=True,
|
||||
show_default=True,
|
||||
default=False,
|
||||
help="Optional. Whether to enable OpenTelemetry for GKE.",
|
||||
)
|
||||
@click.option(
|
||||
"--with_ui",
|
||||
is_flag=True,
|
||||
@@ -1855,6 +1913,7 @@ def cli_deploy_gke(
|
||||
temp_folder: str,
|
||||
port: int,
|
||||
trace_to_cloud: bool,
|
||||
otel_to_cloud: bool,
|
||||
with_ui: bool,
|
||||
adk_version: str,
|
||||
log_level: Optional[str] = None,
|
||||
@@ -1884,6 +1943,7 @@ def cli_deploy_gke(
|
||||
temp_folder=temp_folder,
|
||||
port=port,
|
||||
trace_to_cloud=trace_to_cloud,
|
||||
otel_to_cloud=otel_to_cloud,
|
||||
with_ui=with_ui,
|
||||
log_level=log_level,
|
||||
adk_version=adk_version,
|
||||
|
||||
@@ -46,6 +46,7 @@ from .eval_metrics import EvalMetric
|
||||
from .eval_metrics import EvalMetricResult
|
||||
from .eval_metrics import EvalMetricResultDetails
|
||||
from .eval_metrics import EvalMetricResultPerInvocation
|
||||
from .eval_metrics import Rubric
|
||||
from .eval_result import EvalCaseResult
|
||||
from .eval_set import EvalCase
|
||||
from .eval_set_results_manager import EvalSetResultsManager
|
||||
@@ -67,6 +68,46 @@ def _get_session_id() -> str:
|
||||
return f'{EVAL_SESSION_ID_PREFIX}{str(uuid.uuid4())}'
|
||||
|
||||
|
||||
def _add_rubrics_to_invocation(
|
||||
invocation: Invocation, rubrics_to_add: list[Rubric]
|
||||
):
|
||||
"""Adds rubrics to invocation, throwing ValueError on duplicate rubric_id."""
|
||||
if not invocation.rubrics:
|
||||
invocation.rubrics = []
|
||||
existing_ids = {r.rubric_id for r in invocation.rubrics}
|
||||
for rubric in rubrics_to_add:
|
||||
if rubric.rubric_id in existing_ids:
|
||||
raise ValueError(
|
||||
f"Rubric with rubric_id '{rubric.rubric_id}' already exists."
|
||||
)
|
||||
invocation.rubrics.append(rubric)
|
||||
existing_ids.add(rubric.rubric_id)
|
||||
|
||||
|
||||
def _copy_eval_case_rubrics_to_actual_invocations(
|
||||
eval_case: EvalCase, actual_invocations: list[Invocation]
|
||||
):
|
||||
"""Copies EvalCase level rubrics to all actual invocations."""
|
||||
if hasattr(eval_case, 'rubrics') and eval_case.rubrics:
|
||||
for invocation in actual_invocations:
|
||||
_add_rubrics_to_invocation(invocation, eval_case.rubrics)
|
||||
|
||||
|
||||
def _copy_invocation_rubrics_to_actual_invocations(
|
||||
expected_invocations: Optional[list[Invocation]],
|
||||
actual_invocations: list[Invocation],
|
||||
):
|
||||
"""Copies invocation level rubrics to corresponding actual invocations."""
|
||||
if expected_invocations:
|
||||
for actual_invocation, expected_invocation in zip(
|
||||
actual_invocations, expected_invocations
|
||||
):
|
||||
if expected_invocation.rubrics:
|
||||
_add_rubrics_to_invocation(
|
||||
actual_invocation, expected_invocation.rubrics
|
||||
)
|
||||
|
||||
|
||||
@experimental
|
||||
class LocalEvalService(BaseEvalService):
|
||||
"""An implementation of BaseEvalService, that runs the evals locally."""
|
||||
@@ -249,76 +290,27 @@ class LocalEvalService(BaseEvalService):
|
||||
)
|
||||
)
|
||||
|
||||
actual_invocations = inference_result.inferences
|
||||
expected_invocations = eval_case.conversation
|
||||
|
||||
# 1. Copy EvalCase level rubrics to all actual invocations.
|
||||
_copy_eval_case_rubrics_to_actual_invocations(eval_case, actual_invocations)
|
||||
|
||||
# 2. If expected invocations are present, copy invocation level
|
||||
# rubrics to corresponding actual invocations.
|
||||
_copy_invocation_rubrics_to_actual_invocations(
|
||||
expected_invocations, actual_invocations
|
||||
)
|
||||
|
||||
for eval_metric in evaluate_config.eval_metrics:
|
||||
# Perform evaluation of the metric.
|
||||
try:
|
||||
with client_label_context(EVAL_CLIENT_LABEL):
|
||||
evaluation_result = await self._evaluate_metric(
|
||||
eval_metric=eval_metric,
|
||||
actual_invocations=inference_result.inferences,
|
||||
expected_invocations=eval_case.conversation,
|
||||
conversation_scenario=eval_case.conversation_scenario,
|
||||
)
|
||||
except Exception as e:
|
||||
# We intentionally catch the Exception as we don't want failures to
|
||||
# affect other metric evaluation.
|
||||
logger.error(
|
||||
"Metric evaluation failed for metric `%s` for eval case id '%s'"
|
||||
' with following error `%s`',
|
||||
eval_metric.metric_name,
|
||||
eval_case.eval_id,
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
# We use an empty result.
|
||||
evaluation_result = EvaluationResult(
|
||||
overall_eval_status=EvalStatus.NOT_EVALUATED
|
||||
)
|
||||
|
||||
# Track overall score across all invocations.
|
||||
eval_metric_result_details = EvalMetricResultDetails(
|
||||
rubric_scores=evaluation_result.overall_rubric_scores
|
||||
await self._evaluate_metric_for_eval_case(
|
||||
eval_metric,
|
||||
eval_case,
|
||||
inference_result,
|
||||
eval_metric_result_per_invocation,
|
||||
overall_eval_metric_results,
|
||||
)
|
||||
overall_eval_metric_results.append(
|
||||
EvalMetricResult(
|
||||
score=evaluation_result.overall_score,
|
||||
eval_status=evaluation_result.overall_eval_status,
|
||||
details=eval_metric_result_details,
|
||||
**eval_metric.model_dump(),
|
||||
)
|
||||
)
|
||||
|
||||
if (
|
||||
evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
|
||||
and len(evaluation_result.per_invocation_results)
|
||||
!= len(eval_metric_result_per_invocation)
|
||||
):
|
||||
raise ValueError(
|
||||
'Eval metric should return results for each invocation. Found '
|
||||
f'{len(evaluation_result.per_invocation_results)} results for '
|
||||
f'{len(eval_metric_result_per_invocation)} invocations.'
|
||||
)
|
||||
|
||||
# Track score across individual invocations.
|
||||
for idx, invocation in enumerate(eval_metric_result_per_invocation):
|
||||
invocation_result = (
|
||||
evaluation_result.per_invocation_results[idx]
|
||||
if evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
|
||||
else PerInvocationResult(
|
||||
actual_invocation=invocation.actual_invocation
|
||||
)
|
||||
)
|
||||
eval_metric_result_details = EvalMetricResultDetails(
|
||||
rubric_scores=invocation_result.rubric_scores
|
||||
)
|
||||
invocation.eval_metric_results.append(
|
||||
EvalMetricResult(
|
||||
score=invocation_result.score,
|
||||
eval_status=invocation_result.eval_status,
|
||||
details=eval_metric_result_details,
|
||||
**eval_metric.model_dump(),
|
||||
)
|
||||
)
|
||||
|
||||
final_eval_status = self._generate_final_eval_status(
|
||||
overall_eval_metric_results
|
||||
@@ -342,6 +334,84 @@ class LocalEvalService(BaseEvalService):
|
||||
|
||||
return (inference_result, eval_case_result)
|
||||
|
||||
async def _evaluate_metric_for_eval_case(
|
||||
self,
|
||||
eval_metric: EvalMetric,
|
||||
eval_case: EvalCase,
|
||||
inference_result: InferenceResult,
|
||||
eval_metric_result_per_invocation: list[EvalMetricResultPerInvocation],
|
||||
overall_eval_metric_results: list[EvalMetricResult],
|
||||
):
|
||||
"""Performs evaluation of a metric for a given eval case and inference result."""
|
||||
try:
|
||||
with client_label_context(EVAL_CLIENT_LABEL):
|
||||
evaluation_result = await self._evaluate_metric(
|
||||
eval_metric=eval_metric,
|
||||
actual_invocations=inference_result.inferences,
|
||||
expected_invocations=eval_case.conversation,
|
||||
conversation_scenario=eval_case.conversation_scenario,
|
||||
)
|
||||
except Exception as e:
|
||||
# We intentionally catch the Exception as we don't want failures to
|
||||
# affect other metric evaluation.
|
||||
logger.error(
|
||||
"Metric evaluation failed for metric `%s` for eval case id '%s'"
|
||||
' with following error `%s`',
|
||||
eval_metric.metric_name,
|
||||
eval_case.eval_id,
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
# We use an empty result.
|
||||
evaluation_result = EvaluationResult(
|
||||
overall_eval_status=EvalStatus.NOT_EVALUATED
|
||||
)
|
||||
|
||||
# Track overall score across all invocations.
|
||||
eval_metric_result_details = EvalMetricResultDetails(
|
||||
rubric_scores=evaluation_result.overall_rubric_scores
|
||||
)
|
||||
overall_eval_metric_results.append(
|
||||
EvalMetricResult(
|
||||
score=evaluation_result.overall_score,
|
||||
eval_status=evaluation_result.overall_eval_status,
|
||||
details=eval_metric_result_details,
|
||||
**eval_metric.model_dump(),
|
||||
)
|
||||
)
|
||||
|
||||
if (
|
||||
evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
|
||||
and len(evaluation_result.per_invocation_results)
|
||||
!= len(eval_metric_result_per_invocation)
|
||||
):
|
||||
raise ValueError(
|
||||
'Eval metric should return results for each invocation. Found '
|
||||
f'{len(evaluation_result.per_invocation_results)} results for '
|
||||
f'{len(eval_metric_result_per_invocation)} invocations.'
|
||||
)
|
||||
|
||||
# Track score across individual invocations.
|
||||
for idx, invocation in enumerate(eval_metric_result_per_invocation):
|
||||
invocation_result = (
|
||||
evaluation_result.per_invocation_results[idx]
|
||||
if evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
|
||||
else PerInvocationResult(
|
||||
actual_invocation=invocation.actual_invocation
|
||||
)
|
||||
)
|
||||
eval_metric_result_details = EvalMetricResultDetails(
|
||||
rubric_scores=invocation_result.rubric_scores
|
||||
)
|
||||
invocation.eval_metric_results.append(
|
||||
EvalMetricResult(
|
||||
score=invocation_result.score,
|
||||
eval_status=invocation_result.eval_status,
|
||||
details=eval_metric_result_details,
|
||||
**eval_metric.model_dump(),
|
||||
)
|
||||
)
|
||||
|
||||
async def _evaluate_metric(
|
||||
self,
|
||||
eval_metric: EvalMetric,
|
||||
|
||||
@@ -301,6 +301,7 @@ class RubricBasedEvaluator(LlmAsJudge):
|
||||
invocation_results_summarizer: InvocationResultsSummarizer = (
|
||||
MeanInvocationResultsSummarizer()
|
||||
),
|
||||
rubric_type: Optional[str] = None,
|
||||
):
|
||||
"""Initializes the RubricBasedEvaluator.
|
||||
|
||||
@@ -315,11 +316,14 @@ class RubricBasedEvaluator(LlmAsJudge):
|
||||
to account for the unreliability of the LLM.
|
||||
invocation_results_summarizer: An object that summarizes the results of
|
||||
all invocations in an eval case into a single result.
|
||||
rubric_type: Invocation and case level rubrics will be filtered by this
|
||||
type.
|
||||
"""
|
||||
super().__init__(
|
||||
eval_metric,
|
||||
criterion_type=criterion_type,
|
||||
)
|
||||
self._rubric_type = rubric_type
|
||||
self._auto_rater_prompt_template = ""
|
||||
self._auto_rater_response_parser = auto_rater_response_parser
|
||||
self._per_invocation_results_aggregator = per_invocation_results_aggregator
|
||||
@@ -328,28 +332,72 @@ class RubricBasedEvaluator(LlmAsJudge):
|
||||
assert self._criterion.rubrics, "Rubrics are required."
|
||||
|
||||
self._rubrics: list[Rubric] = self._criterion.rubrics
|
||||
self._effective_rubrics_list: Optional[list[Rubric]] = None
|
||||
|
||||
self._normalized_rubric_to_id_map = {
|
||||
_normalize_text(r.rubric_content.text_property): r.rubric_id
|
||||
for r in self._rubrics
|
||||
}
|
||||
|
||||
def create_effective_rubrics_list(
|
||||
self,
|
||||
invocation_rubrics: Optional[list[Rubric]],
|
||||
) -> None:
|
||||
rubrics_by_id = {}
|
||||
|
||||
def _add_rubrics(rubrics_to_add: list[Rubric], scope_name: str):
|
||||
for r in rubrics_to_add:
|
||||
if r.rubric_id in rubrics_by_id:
|
||||
raise ValueError(
|
||||
f"Rubric with rubric_id '{r.rubric_id}' already exists. Rubric"
|
||||
f" defined in {scope_name} conflicts with an existing rubric."
|
||||
)
|
||||
rubrics_by_id[r.rubric_id] = r
|
||||
|
||||
_add_rubrics(self._rubrics, "criterion")
|
||||
|
||||
if invocation_rubrics:
|
||||
filtered_invocation_rubrics = invocation_rubrics
|
||||
if self._rubric_type:
|
||||
filtered_invocation_rubrics = [
|
||||
r for r in invocation_rubrics if r.type == self._rubric_type
|
||||
]
|
||||
_add_rubrics(filtered_invocation_rubrics, "invocation")
|
||||
|
||||
self._effective_rubrics_list = list(rubrics_by_id.values())
|
||||
|
||||
def get_effective_rubrics_list(self) -> list[Rubric]:
|
||||
"""Returns the effective rubrics list."""
|
||||
if self._effective_rubrics_list is None:
|
||||
raise ValueError(
|
||||
"Effective rubrics list not initialized. Call"
|
||||
" create_effective_rubrics_list() first."
|
||||
)
|
||||
return self._effective_rubrics_list
|
||||
|
||||
@override
|
||||
def convert_auto_rater_response_to_score(
|
||||
self, auto_rater_response: LlmResponse
|
||||
self,
|
||||
auto_rater_response: LlmResponse,
|
||||
) -> AutoRaterScore:
|
||||
"""Returns an AutoRaterScore generated from AutoRater's response."""
|
||||
response_text = get_text_from_content(auto_rater_response.content)
|
||||
rubric_responses = self._auto_rater_response_parser.parse(response_text)
|
||||
rubric_scores = []
|
||||
|
||||
normalized_rubric_to_rubric_map = {}
|
||||
for r in self.get_effective_rubrics_list():
|
||||
normalized_rubric_to_rubric_map[
|
||||
_normalize_text(r.rubric_content.text_property)
|
||||
] = r
|
||||
|
||||
for rubric_response in rubric_responses:
|
||||
normalized_rubric = _normalize_text(rubric_response.property_text)
|
||||
rubric_id = self._normalized_rubric_to_id_map.get(normalized_rubric, None)
|
||||
if rubric_id:
|
||||
normalized_rubric_text = _normalize_text(rubric_response.property_text)
|
||||
rubric = normalized_rubric_to_rubric_map.get(normalized_rubric_text, None)
|
||||
if rubric:
|
||||
rubric_scores.append(
|
||||
RubricScore(
|
||||
rubric_id=rubric_id,
|
||||
rubric_id=rubric.rubric_id,
|
||||
rationale=rubric_response.rationale,
|
||||
score=rubric_response.score,
|
||||
)
|
||||
|
||||
@@ -25,6 +25,7 @@ from .eval_case import Invocation
|
||||
from .eval_case import InvocationEvents
|
||||
from .eval_metrics import EvalMetric
|
||||
from .eval_metrics import RubricsBasedCriterion
|
||||
from .eval_rubrics import Rubric
|
||||
from .llm_as_judge_utils import get_text_from_content
|
||||
from .llm_as_judge_utils import get_tool_calls_and_responses_as_json_str
|
||||
from .llm_as_judge_utils import get_tool_declarations_as_json_str
|
||||
@@ -252,11 +253,13 @@ class RubricBasedFinalResponseQualityV1Evaluator(RubricBasedEvaluator):
|
||||
"""
|
||||
|
||||
criterion_type: ClassVar[type[RubricsBasedCriterion]] = RubricsBasedCriterion
|
||||
RUBRIC_TYPE: ClassVar[str] = "FINAL_RESPONSE_QUALITY"
|
||||
|
||||
def __init__(self, eval_metric: EvalMetric):
|
||||
super().__init__(
|
||||
eval_metric,
|
||||
criterion_type=RubricBasedFinalResponseQualityV1Evaluator.criterion_type,
|
||||
rubric_type=RubricBasedFinalResponseQualityV1Evaluator.RUBRIC_TYPE,
|
||||
)
|
||||
self._auto_rater_prompt_template = (
|
||||
_RUBRIC_BASED_FINAL_RESPONSE_QUALITY_V1_PROMPT
|
||||
@@ -264,15 +267,19 @@ class RubricBasedFinalResponseQualityV1Evaluator(RubricBasedEvaluator):
|
||||
|
||||
@override
|
||||
def format_auto_rater_prompt(
|
||||
self, actual_invocation: Invocation, _: Optional[Invocation]
|
||||
self,
|
||||
actual_invocation: Invocation,
|
||||
_: Optional[Invocation],
|
||||
) -> str:
|
||||
"""Returns the autorater prompt."""
|
||||
|
||||
self.create_effective_rubrics_list(actual_invocation.rubrics)
|
||||
user_input = get_text_from_content(actual_invocation.user_content)
|
||||
final_response = get_text_from_content(actual_invocation.final_response)
|
||||
rubrics = "\n* ".join(
|
||||
[r.rubric_content.text_property for r in self._rubrics]
|
||||
)
|
||||
|
||||
rubrics_text = "\n".join([
|
||||
f"* {r.rubric_content.text_property}"
|
||||
for r in self._effective_rubrics_list
|
||||
])
|
||||
|
||||
developer_instructions = ""
|
||||
tool_declarations = "Agent has no tools."
|
||||
@@ -299,7 +306,7 @@ class RubricBasedFinalResponseQualityV1Evaluator(RubricBasedEvaluator):
|
||||
user_input=user_input,
|
||||
response_steps=response_steps,
|
||||
final_response=final_response,
|
||||
rubrics=rubrics,
|
||||
rubrics=rubrics_text,
|
||||
)
|
||||
|
||||
return auto_rater_prompt
|
||||
|
||||
@@ -154,27 +154,33 @@ class RubricBasedToolUseV1Evaluator(RubricBasedEvaluator):
|
||||
"""
|
||||
|
||||
criterion_type: ClassVar[type[RubricsBasedCriterion]] = RubricsBasedCriterion
|
||||
RUBRIC_TYPE: ClassVar[str] = "TOOL_USE_QUALITY"
|
||||
|
||||
def __init__(self, eval_metric: EvalMetric):
|
||||
super().__init__(
|
||||
eval_metric,
|
||||
criterion_type=RubricBasedToolUseV1Evaluator.criterion_type,
|
||||
rubric_type=RubricBasedToolUseV1Evaluator.RUBRIC_TYPE,
|
||||
)
|
||||
self._auto_rater_prompt_template = _RUBRIC_BASED_TOOL_USE_QUALITY_V1_PROMPT
|
||||
|
||||
@override
|
||||
def format_auto_rater_prompt(
|
||||
self, actual_invocation: Invocation, _: Optional[Invocation]
|
||||
self,
|
||||
actual_invocation: Invocation,
|
||||
_: Optional[Invocation],
|
||||
) -> str:
|
||||
"""Returns the autorater prompt."""
|
||||
|
||||
self.create_effective_rubrics_list(actual_invocation.rubrics)
|
||||
user_input = get_text_from_content(actual_invocation.user_content)
|
||||
tool_usage = get_tool_calls_and_responses_as_json_str(
|
||||
actual_invocation.intermediate_data
|
||||
)
|
||||
rubrics = "\n* ".join(
|
||||
[r.rubric_content.text_property for r in self._rubrics]
|
||||
)
|
||||
|
||||
rubrics_text = "\n".join([
|
||||
f"* {r.rubric_content.text_property}"
|
||||
for r in self._effective_rubrics_list
|
||||
])
|
||||
|
||||
app_details = actual_invocation.app_details
|
||||
tool_declarations = "Agent has no tools."
|
||||
@@ -185,5 +191,5 @@ class RubricBasedToolUseV1Evaluator(RubricBasedEvaluator):
|
||||
tool_declarations=tool_declarations,
|
||||
user_input=user_input,
|
||||
tool_usage=tool_usage,
|
||||
rubrics=rubrics,
|
||||
rubrics=rubrics_text,
|
||||
)
|
||||
|
||||
@@ -14,8 +14,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Generator
|
||||
import warnings
|
||||
|
||||
from ..utils.env_utils import is_env_enabled
|
||||
@@ -24,17 +26,24 @@ from ..utils.env_utils import is_env_enabled
|
||||
class FeatureName(str, Enum):
|
||||
"""Feature names."""
|
||||
|
||||
AUTHENTICATED_FUNCTION_TOOL = "AUTHENTICATED_FUNCTION_TOOL"
|
||||
BASE_AUTHENTICATED_TOOL = "BASE_AUTHENTICATED_TOOL"
|
||||
BIG_QUERY_TOOLSET = "BIG_QUERY_TOOLSET"
|
||||
BIG_QUERY_TOOL_CONFIG = "BIG_QUERY_TOOL_CONFIG"
|
||||
BIGTABLE_TOOL_SETTINGS = "BIGTABLE_TOOL_SETTINGS"
|
||||
BIGTABLE_TOOLSET = "BIGTABLE_TOOLSET"
|
||||
COMPUTER_USE = "COMPUTER_USE"
|
||||
GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
|
||||
GOOGLE_TOOL = "GOOGLE_TOOL"
|
||||
JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
|
||||
PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
|
||||
PUBSUB_TOOL_CONFIG = "PUBSUB_TOOL_CONFIG"
|
||||
PUBSUB_TOOLSET = "PUBSUB_TOOLSET"
|
||||
SPANNER_TOOLSET = "SPANNER_TOOLSET"
|
||||
SPANNER_TOOL_SETTINGS = "SPANNER_TOOL_SETTINGS"
|
||||
SPANNER_VECTOR_STORE = "SPANNER_VECTOR_STORE"
|
||||
TOOL_CONFIG = "TOOL_CONFIG"
|
||||
TOOL_CONFIRMATION = "TOOL_CONFIRMATION"
|
||||
|
||||
|
||||
class FeatureStage(Enum):
|
||||
@@ -67,6 +76,12 @@ class FeatureConfig:
|
||||
|
||||
# Central registry: FeatureName -> FeatureConfig
|
||||
_FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
|
||||
FeatureName.AUTHENTICATED_FUNCTION_TOOL: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.BASE_AUTHENTICATED_TOOL: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.BIG_QUERY_TOOLSET: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
@@ -76,6 +91,9 @@ _FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
|
||||
FeatureName.BIGTABLE_TOOL_SETTINGS: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.BIGTABLE_TOOLSET: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.COMPUTER_USE: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
@@ -91,6 +109,9 @@ _FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
|
||||
FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.PUBSUB_TOOL_CONFIG: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.PUBSUB_TOOLSET: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
@@ -100,6 +121,15 @@ _FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
|
||||
FeatureName.SPANNER_TOOL_SETTINGS: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.SPANNER_VECTOR_STORE: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.TOOL_CONFIG: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
FeatureName.TOOL_CONFIRMATION: FeatureConfig(
|
||||
FeatureStage.EXPERIMENTAL, default_on=True
|
||||
),
|
||||
}
|
||||
|
||||
# Track which experimental features have already warned (warn only once)
|
||||
@@ -240,3 +270,52 @@ def _emit_non_stable_warning_once(
|
||||
f"[{feature_stage.name.upper()}] feature {feature_name} is enabled."
|
||||
)
|
||||
warnings.warn(full_message, category=UserWarning, stacklevel=4)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def temporary_feature_override(
|
||||
feature_name: FeatureName,
|
||||
enabled: bool,
|
||||
) -> Generator[None, None, None]:
|
||||
"""Temporarily override a feature's enabled state within a context.
|
||||
|
||||
This context manager is useful for testing or temporarily enabling/disabling
|
||||
a feature within a specific scope. The original state is restored when the
|
||||
context exits.
|
||||
|
||||
Args:
|
||||
feature_name: The feature name to override.
|
||||
enabled: Whether the feature should be enabled.
|
||||
|
||||
Yields:
|
||||
None
|
||||
|
||||
Example:
|
||||
```python
|
||||
from google.adk.features import FeatureName, temporary_feature_override
|
||||
|
||||
# Temporarily enable a feature for testing
|
||||
with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True):
|
||||
# Feature is enabled here
|
||||
result = some_function_that_checks_feature()
|
||||
# Feature is restored to original state here
|
||||
```
|
||||
"""
|
||||
config = _get_feature_config(feature_name)
|
||||
if config is None:
|
||||
raise ValueError(f"Feature {feature_name} is not registered.")
|
||||
|
||||
# Save the original override state
|
||||
had_override = feature_name in _FEATURE_OVERRIDES
|
||||
original_value = _FEATURE_OVERRIDES.get(feature_name)
|
||||
|
||||
# Apply the temporary override
|
||||
_FEATURE_OVERRIDES[feature_name] = enabled
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Restore the original state
|
||||
if had_override:
|
||||
_FEATURE_OVERRIDES[feature_name] = original_value
|
||||
else:
|
||||
_FEATURE_OVERRIDES.pop(feature_name, None)
|
||||
|
||||
@@ -116,6 +116,10 @@ class BaseLlmFlow(ABC):
|
||||
attempt += 1
|
||||
if not llm_request.live_connect_config:
|
||||
llm_request.live_connect_config = types.LiveConnectConfig()
|
||||
if not llm_request.live_connect_config.session_resumption:
|
||||
llm_request.live_connect_config.session_resumption = (
|
||||
types.SessionResumptionConfig()
|
||||
)
|
||||
llm_request.live_connect_config.session_resumption.handle = (
|
||||
invocation_context.live_session_resumption_handle
|
||||
)
|
||||
@@ -299,7 +303,6 @@ class BaseLlmFlow(ABC):
|
||||
else:
|
||||
return invocation_context.agent.name
|
||||
|
||||
assert invocation_context.live_request_queue
|
||||
try:
|
||||
while True:
|
||||
async with Aclosing(llm_connection.receive()) as agen:
|
||||
|
||||
@@ -41,11 +41,13 @@ class GeminiLlmConnection(BaseLlmConnection):
|
||||
self,
|
||||
gemini_session: live.AsyncSession,
|
||||
api_backend: GoogleLLMVariant = GoogleLLMVariant.VERTEX_AI,
|
||||
model_version: str | None = None,
|
||||
):
|
||||
self._gemini_session = gemini_session
|
||||
self._input_transcription_text: str = ''
|
||||
self._output_transcription_text: str = ''
|
||||
self._api_backend = api_backend
|
||||
self._model_version = model_version
|
||||
|
||||
async def send_history(self, history: list[types.Content]):
|
||||
"""Sends the conversation history to the gemini model.
|
||||
@@ -162,7 +164,11 @@ class GeminiLlmConnection(BaseLlmConnection):
|
||||
async for message in agen:
|
||||
logger.debug('Got LLM Live message: %s', message)
|
||||
if message.usage_metadata:
|
||||
yield LlmResponse(usage_metadata=message.usage_metadata)
|
||||
# Tracks token usage data per model.
|
||||
yield LlmResponse(
|
||||
usage_metadata=message.usage_metadata,
|
||||
model_version=self._model_version,
|
||||
)
|
||||
if message.server_content:
|
||||
content = message.server_content.model_turn
|
||||
if content and content.parts:
|
||||
|
||||
@@ -402,7 +402,11 @@ class Gemini(BaseLlm):
|
||||
async with self._live_api_client.aio.live.connect(
|
||||
model=llm_request.model, config=llm_request.live_connect_config
|
||||
) as live_session:
|
||||
yield GeminiLlmConnection(live_session, api_backend=self._api_backend)
|
||||
yield GeminiLlmConnection(
|
||||
live_session,
|
||||
api_backend=self._api_backend,
|
||||
model_version=llm_request.model,
|
||||
)
|
||||
|
||||
async def _adapt_computer_use_tool(self, llm_request: LlmRequest) -> None:
|
||||
"""Adapt the google computer use predefined functions to the adk computer use toolset."""
|
||||
|
||||
@@ -181,6 +181,45 @@ def _infer_mime_type_from_uri(uri: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def _looks_like_openai_file_id(file_uri: str) -> bool:
|
||||
"""Returns True when file_uri resembles an OpenAI/Azure file id."""
|
||||
return file_uri.startswith("file-")
|
||||
|
||||
|
||||
def _redact_file_uri_for_log(
|
||||
file_uri: str, *, display_name: str | None = None
|
||||
) -> str:
|
||||
"""Returns a privacy-preserving identifier for logs."""
|
||||
if display_name:
|
||||
return display_name
|
||||
if _looks_like_openai_file_id(file_uri):
|
||||
return "file-<redacted>"
|
||||
try:
|
||||
parsed = urlparse(file_uri)
|
||||
except ValueError:
|
||||
return "<unparseable>"
|
||||
if not parsed.scheme:
|
||||
return "<unknown>"
|
||||
segments = [segment for segment in parsed.path.split("/") if segment]
|
||||
tail = segments[-1] if segments else ""
|
||||
if tail:
|
||||
return f"{parsed.scheme}://<redacted>/{tail}"
|
||||
return f"{parsed.scheme}://<redacted>"
|
||||
|
||||
|
||||
def _requires_file_uri_fallback(
|
||||
provider: str, model: str, file_uri: str
|
||||
) -> bool:
|
||||
"""Returns True when `file_uri` should not be sent as a file content block."""
|
||||
if provider in _FILE_ID_REQUIRED_PROVIDERS:
|
||||
return not _looks_like_openai_file_id(file_uri)
|
||||
if provider == "anthropic":
|
||||
return True
|
||||
if provider == "vertex_ai" and not _is_litellm_gemini_model(model):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _decode_inline_text_data(raw_bytes: bytes) -> str:
|
||||
"""Decodes inline file bytes that represent textual content."""
|
||||
try:
|
||||
@@ -447,6 +486,7 @@ async def _content_to_message_param(
|
||||
content: types.Content,
|
||||
*,
|
||||
provider: str = "",
|
||||
model: str = "",
|
||||
) -> Union[Message, list[Message]]:
|
||||
"""Converts a types.Content to a litellm Message or list of Messages.
|
||||
|
||||
@@ -456,12 +496,14 @@ async def _content_to_message_param(
|
||||
Args:
|
||||
content: The content to convert.
|
||||
provider: The LLM provider name (e.g., "openai", "azure").
|
||||
model: The LiteLLM model string, used for provider-specific behavior.
|
||||
|
||||
Returns:
|
||||
A litellm Message, a list of litellm Messages.
|
||||
"""
|
||||
|
||||
tool_messages = []
|
||||
tool_messages: list[Message] = []
|
||||
non_tool_parts: list[types.Part] = []
|
||||
for part in content.parts:
|
||||
if part.function_response:
|
||||
response = part.function_response.response
|
||||
@@ -477,15 +519,30 @@ async def _content_to_message_param(
|
||||
content=response_content,
|
||||
)
|
||||
)
|
||||
if tool_messages:
|
||||
else:
|
||||
non_tool_parts.append(part)
|
||||
|
||||
if tool_messages and not non_tool_parts:
|
||||
return tool_messages if len(tool_messages) > 1 else tool_messages[0]
|
||||
|
||||
if tool_messages and non_tool_parts:
|
||||
follow_up = await _content_to_message_param(
|
||||
types.Content(role=content.role, parts=non_tool_parts),
|
||||
provider=provider,
|
||||
)
|
||||
follow_up_messages = (
|
||||
follow_up if isinstance(follow_up, list) else [follow_up]
|
||||
)
|
||||
return tool_messages + follow_up_messages
|
||||
|
||||
# Handle user or assistant messages
|
||||
role = _to_litellm_role(content.role)
|
||||
|
||||
if role == "user":
|
||||
user_parts = [part for part in content.parts if not part.thought]
|
||||
message_content = await _get_content(user_parts, provider=provider) or None
|
||||
message_content = (
|
||||
await _get_content(user_parts, provider=provider, model=model) or None
|
||||
)
|
||||
return ChatCompletionUserMessage(role="user", content=message_content)
|
||||
else: # assistant/model
|
||||
tool_calls = []
|
||||
@@ -509,7 +566,7 @@ async def _content_to_message_param(
|
||||
content_parts.append(part)
|
||||
|
||||
final_content = (
|
||||
await _get_content(content_parts, provider=provider)
|
||||
await _get_content(content_parts, provider=provider, model=model)
|
||||
if content_parts
|
||||
else None
|
||||
)
|
||||
@@ -606,6 +663,7 @@ async def _get_content(
|
||||
parts: Iterable[types.Part],
|
||||
*,
|
||||
provider: str = "",
|
||||
model: str = "",
|
||||
) -> OpenAIMessageContent:
|
||||
"""Converts a list of parts to litellm content.
|
||||
|
||||
@@ -615,6 +673,8 @@ async def _get_content(
|
||||
Args:
|
||||
parts: The parts to convert.
|
||||
provider: The LLM provider name (e.g., "openai", "azure").
|
||||
model: The LiteLLM model string (e.g., "openai/gpt-4o",
|
||||
"vertex_ai/gemini-2.5-flash").
|
||||
|
||||
Returns:
|
||||
The litellm content.
|
||||
@@ -695,6 +755,32 @@ async def _get_content(
|
||||
f"{part.inline_data.mime_type}."
|
||||
)
|
||||
elif part.file_data and part.file_data.file_uri:
|
||||
if (
|
||||
provider in _FILE_ID_REQUIRED_PROVIDERS
|
||||
and _looks_like_openai_file_id(part.file_data.file_uri)
|
||||
):
|
||||
content_objects.append({
|
||||
"type": "file",
|
||||
"file": {"file_id": part.file_data.file_uri},
|
||||
})
|
||||
continue
|
||||
|
||||
if _requires_file_uri_fallback(provider, model, part.file_data.file_uri):
|
||||
logger.debug(
|
||||
"File URI %s not supported for provider %s, using text fallback",
|
||||
_redact_file_uri_for_log(
|
||||
part.file_data.file_uri,
|
||||
display_name=part.file_data.display_name,
|
||||
),
|
||||
provider,
|
||||
)
|
||||
identifier = part.file_data.display_name or part.file_data.file_uri
|
||||
content_objects.append({
|
||||
"type": "text",
|
||||
"text": f'[File reference: "{identifier}"]',
|
||||
})
|
||||
continue
|
||||
|
||||
file_object: ChatCompletionFileUrlObject = {
|
||||
"file_id": part.file_data.file_uri,
|
||||
}
|
||||
@@ -1349,7 +1435,7 @@ async def _get_completion_inputs(
|
||||
messages: List[Message] = []
|
||||
for content in llm_request.contents or []:
|
||||
message_param_or_list = await _content_to_message_param(
|
||||
content, provider=provider
|
||||
content, provider=provider, model=model
|
||||
)
|
||||
if isinstance(message_param_or_list, list):
|
||||
messages.extend(message_param_or_list)
|
||||
|
||||
@@ -36,6 +36,7 @@ from typing import TYPE_CHECKING
|
||||
import uuid
|
||||
import weakref
|
||||
|
||||
from google.api_core import client_options
|
||||
from google.api_core.exceptions import InternalServerError
|
||||
from google.api_core.exceptions import ServiceUnavailable
|
||||
from google.api_core.exceptions import TooManyRequests
|
||||
@@ -1352,19 +1353,31 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
|
||||
if _GLOBAL_WRITE_CLIENT is None:
|
||||
|
||||
def get_credentials():
|
||||
creds, _ = google.auth.default(
|
||||
creds, project_id = google.auth.default(
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
||||
)
|
||||
return creds
|
||||
return creds, project_id
|
||||
|
||||
creds = await loop.run_in_executor(self._executor, get_credentials)
|
||||
creds, project_id = await loop.run_in_executor(
|
||||
self._executor, get_credentials
|
||||
)
|
||||
quota_project_id = (
|
||||
getattr(creds, "quota_project_id", None) or project_id
|
||||
)
|
||||
options = (
|
||||
client_options.ClientOptions(quota_project_id=quota_project_id)
|
||||
if quota_project_id
|
||||
else None
|
||||
)
|
||||
client_info = gapic_client_info.ClientInfo(
|
||||
user_agent=f"google-adk-bq-logger/{__version__}"
|
||||
)
|
||||
# Initialize the async client in the current event loop, not in the
|
||||
# executor.
|
||||
_GLOBAL_WRITE_CLIENT = BigQueryWriteAsyncClient(
|
||||
credentials=creds, client_info=client_info
|
||||
credentials=creds,
|
||||
client_info=client_info,
|
||||
client_options=options,
|
||||
)
|
||||
self.write_client = _GLOBAL_WRITE_CLIENT
|
||||
|
||||
|
||||
+26
-17
@@ -78,7 +78,9 @@ def _is_transcription(event: Event) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _has_non_empty_transcription_text(transcription) -> bool:
|
||||
def _has_non_empty_transcription_text(
|
||||
transcription: types.Transcription,
|
||||
) -> bool:
|
||||
return bool(
|
||||
transcription and transcription.text and transcription.text.strip()
|
||||
)
|
||||
@@ -151,15 +153,21 @@ class Runner:
|
||||
"""Initializes the Runner.
|
||||
|
||||
Developers should provide either an `app` instance or both `app_name` and
|
||||
`agent`. Providing a mix of `app` and `app_name`/`agent` will result in a
|
||||
`ValueError`. Providing `app` is the recommended way to create a runner.
|
||||
`agent`. When `app` is provided, `app_name` can optionally override the
|
||||
app's name (useful for deployment scenarios like Agent Engine where the
|
||||
resource name differs from the app's identifier). However, `agent` should
|
||||
not be provided when `app` is provided. Providing `app` is the recommended
|
||||
way to create a runner.
|
||||
|
||||
Args:
|
||||
app: An optional `App` instance. If provided, `app_name` and `agent`
|
||||
should not be specified.
|
||||
app: An optional `App` instance. If provided, `agent` should not be
|
||||
specified. `app_name` can optionally override `app.name`.
|
||||
app_name: The application name of the runner. Required if `app` is not
|
||||
provided.
|
||||
agent: The root agent to run. Required if `app` is not provided.
|
||||
provided. If `app` is provided, this can optionally override `app.name`
|
||||
(e.g., for deployment scenarios where a resource name differs from the
|
||||
app identifier).
|
||||
agent: The root agent to run. Required if `app` is not provided. Should
|
||||
not be provided when `app` is provided.
|
||||
plugins: Deprecated. A list of plugins for the runner. Please use the
|
||||
`app` argument to provide plugins instead.
|
||||
artifact_service: The artifact service for the runner.
|
||||
@@ -169,8 +177,8 @@ class Runner:
|
||||
plugin_close_timeout: The timeout in seconds for plugin close methods.
|
||||
|
||||
Raises:
|
||||
ValueError: If `app` is provided along with `app_name` or `plugins`, or
|
||||
if `app` is not provided but either `app_name` or `agent` is missing.
|
||||
ValueError: If `app` is provided along with `agent` or `plugins`, or if
|
||||
`app` is not provided but either `app_name` or `agent` is missing.
|
||||
"""
|
||||
self.app = app
|
||||
(
|
||||
@@ -211,7 +219,8 @@ class Runner:
|
||||
|
||||
Args:
|
||||
app: An optional `App` instance.
|
||||
app_name: The application name of the runner.
|
||||
app_name: The application name of the runner. Can override app.name when
|
||||
app is provided.
|
||||
agent: The root agent to run.
|
||||
plugins: A list of plugins for the runner.
|
||||
|
||||
@@ -230,10 +239,6 @@ class Runner:
|
||||
)
|
||||
|
||||
if app:
|
||||
if app_name:
|
||||
raise ValueError(
|
||||
'When app is provided, app_name should not be provided.'
|
||||
)
|
||||
if agent:
|
||||
raise ValueError('When app is provided, agent should not be provided.')
|
||||
if plugins:
|
||||
@@ -241,7 +246,9 @@ class Runner:
|
||||
'When app is provided, plugins should not be provided and should be'
|
||||
' provided in the app instead.'
|
||||
)
|
||||
app_name = app.name
|
||||
# Allow app_name to override app.name (useful for deployment scenarios
|
||||
# like Agent Engine where resource names differ from app identifiers)
|
||||
app_name = app_name or app.name
|
||||
agent = app.root_agent
|
||||
plugins = app.plugins
|
||||
context_cache_config = app.context_cache_config
|
||||
@@ -950,6 +957,8 @@ class Runner:
|
||||
raise ValueError(
|
||||
'Either session or user_id and session_id must be provided.'
|
||||
)
|
||||
if live_request_queue is None:
|
||||
raise ValueError('live_request_queue is required for run_live.')
|
||||
if session is not None:
|
||||
warnings.warn(
|
||||
'The `session` parameter is deprecated. Please use `user_id` and'
|
||||
@@ -1378,7 +1387,7 @@ class Runner:
|
||||
self,
|
||||
session: Session,
|
||||
*,
|
||||
live_request_queue: Optional[LiveRequestQueue] = None,
|
||||
live_request_queue: LiveRequestQueue,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
) -> InvocationContext:
|
||||
"""Creates a new invocation context for live multi-agent."""
|
||||
@@ -1386,7 +1395,7 @@ class Runner:
|
||||
|
||||
# For live multi-agents system, we need model's text transcription as
|
||||
# context for the transferred agent.
|
||||
if self.agent.sub_agents and live_request_queue:
|
||||
if self.agent.sub_agents:
|
||||
if 'AUDIO' in run_config.response_modalities:
|
||||
if not run_config.output_audio_transcription:
|
||||
run_config.output_audio_transcription = (
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Database schema version check utility."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
@@ -32,8 +33,11 @@ def _get_schema_version_impl(inspector, connection) -> str:
|
||||
"""Gets DB schema version using inspector and connection."""
|
||||
if inspector.has_table("adk_internal_metadata"):
|
||||
try:
|
||||
key_col = inspector.dialect.identifier_preparer.quote("key")
|
||||
result = connection.execute(
|
||||
text("SELECT value FROM adk_internal_metadata WHERE key = :key"),
|
||||
text(
|
||||
f"SELECT value FROM adk_internal_metadata WHERE {key_col} = :key"
|
||||
),
|
||||
{"key": SCHEMA_VERSION_KEY},
|
||||
).fetchone()
|
||||
if result:
|
||||
@@ -49,6 +53,7 @@ def _get_schema_version_impl(inspector, connection) -> str:
|
||||
e,
|
||||
)
|
||||
raise
|
||||
|
||||
# Metadata table doesn't exist, check for v0 schema.
|
||||
# V0 schema has an 'events' table with an 'actions' column.
|
||||
if inspector.has_table("events"):
|
||||
|
||||
@@ -22,6 +22,8 @@ import sqlite3
|
||||
import time
|
||||
from typing import Any
|
||||
from typing import Optional
|
||||
from urllib.parse import unquote
|
||||
from urllib.parse import urlparse
|
||||
import uuid
|
||||
|
||||
import aiosqlite
|
||||
@@ -91,6 +93,42 @@ CREATE_SCHEMA_SQL = "\n".join([
|
||||
])
|
||||
|
||||
|
||||
def _parse_db_path(db_path: str) -> tuple[str, str, bool]:
|
||||
"""Normalizes a SQLite db path from a URL or filesystem path.
|
||||
|
||||
Returns:
|
||||
A tuple of:
|
||||
- filesystem path (for `os.path.exists` and user-facing messages)
|
||||
- value to pass to sqlite/aiosqlite connect
|
||||
- whether to pass `uri=True` to sqlite/aiosqlite connect
|
||||
|
||||
Notes:
|
||||
When a SQLAlchemy-style SQLite URL is provided, this follows SQLAlchemy's
|
||||
conventions:
|
||||
- `sqlite:///relative.db` is a path relative to the current working dir.
|
||||
- `sqlite:////absolute.db` is an absolute filesystem path.
|
||||
"""
|
||||
if not db_path.startswith(("sqlite:", "sqlite+aiosqlite:")):
|
||||
return db_path, db_path, False
|
||||
|
||||
parsed = urlparse(db_path)
|
||||
raw_path = unquote(parsed.path)
|
||||
if not raw_path:
|
||||
return db_path, db_path, False
|
||||
|
||||
normalized_path = raw_path
|
||||
if normalized_path.startswith("//"):
|
||||
normalized_path = normalized_path[1:]
|
||||
elif normalized_path.startswith("/"):
|
||||
normalized_path = normalized_path[1:]
|
||||
|
||||
if parsed.query:
|
||||
# sqlite3 only treats the filename as a URI when it starts with `file:`.
|
||||
return normalized_path, f"file:{normalized_path}?{parsed.query}", True
|
||||
|
||||
return normalized_path, normalized_path, False
|
||||
|
||||
|
||||
class SqliteSessionService(BaseSessionService):
|
||||
"""A session service that uses an SQLite database for storage via aiosqlite.
|
||||
|
||||
@@ -100,17 +138,19 @@ class SqliteSessionService(BaseSessionService):
|
||||
|
||||
def __init__(self, db_path: str):
|
||||
"""Initializes the SQLite session service with a database path."""
|
||||
self._db_path = db_path
|
||||
self._db_path, self._db_connect_path, self._db_connect_uri = _parse_db_path(
|
||||
db_path
|
||||
)
|
||||
|
||||
if self._is_migration_needed():
|
||||
raise RuntimeError(
|
||||
f"Database {db_path} seems to use an old schema."
|
||||
f"Database {self._db_path} seems to use an old schema."
|
||||
" Please run the migration command to"
|
||||
" migrate it to the new schema. Example: `python -m"
|
||||
" google.adk.sessions.migration.migrate_from_sqlalchemy_sqlite"
|
||||
f" --source_db_path {db_path} --dest_db_path"
|
||||
f" {db_path}.new` then backup {db_path} and rename"
|
||||
f" {db_path}.new to {db_path}."
|
||||
f" --source_db_path {self._db_path} --dest_db_path"
|
||||
f" {self._db_path}.new` then backup {self._db_path} and rename"
|
||||
f" {self._db_path}.new to {self._db_path}."
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -415,7 +455,9 @@ class SqliteSessionService(BaseSessionService):
|
||||
@asynccontextmanager
|
||||
async def _get_db_connection(self):
|
||||
"""Connects to the db and performs initial setup."""
|
||||
async with aiosqlite.connect(self._db_path) as db:
|
||||
async with aiosqlite.connect(
|
||||
self._db_connect_path, uri=self._db_connect_uri
|
||||
) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
await db.execute(PRAGMA_FOREIGN_KEYS)
|
||||
await db.executescript(CREATE_SCHEMA_SQL)
|
||||
@@ -514,7 +556,9 @@ class SqliteSessionService(BaseSessionService):
|
||||
if not os.path.exists(self._db_path):
|
||||
return False
|
||||
try:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
with sqlite3.connect(
|
||||
self._db_connect_path, uri=self._db_connect_uri
|
||||
) as conn:
|
||||
cursor = conn.cursor()
|
||||
# Check if events table exists
|
||||
cursor.execute(
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user