diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 7a23a6cc8e..ee7f956571 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -16,7 +16,8 @@ from typing import Any from typing import Optional -import uuid + +from google.adk.platform import uuid as platform_uuid from google.genai import types from pydantic import BaseModel @@ -409,4 +410,4 @@ def _find_matching_function_call( def new_invocation_context_id() -> str: - return "e-" + str(uuid.uuid4()) + return "e-" + platform_uuid.new_uuid() diff --git a/src/google/adk/artifacts/base_artifact_service.py b/src/google/adk/artifacts/base_artifact_service.py index 1a265f8ad9..2421b6e2b1 100644 --- a/src/google/adk/artifacts/base_artifact_service.py +++ b/src/google/adk/artifacts/base_artifact_service.py @@ -19,6 +19,8 @@ from typing import Any from typing import Optional +from google.adk.platform import time as platform_time + from google.genai import types from pydantic import alias_generators from pydantic import BaseModel @@ -47,7 +49,7 @@ class ArtifactVersion(BaseModel): description="Optional user-supplied metadata stored with the artifact.", ) create_time: float = Field( - default_factory=lambda: datetime.now().timestamp(), + default_factory=lambda: platform_time.get_time(), description=( "Unix timestamp (seconds) when the version record was created." ), diff --git a/src/google/adk/events/event.py b/src/google/adk/events/event.py index 2c6a6cd66c..89e2bf05e3 100644 --- a/src/google/adk/events/event.py +++ b/src/google/adk/events/event.py @@ -18,6 +18,8 @@ from typing import Optional import uuid +from google.adk.platform import time as platform_time +from google.adk.platform import uuid as platform_uuid from google.genai import types from pydantic import alias_generators from pydantic import ConfigDict @@ -70,7 +72,7 @@ class Event(LlmResponse): # Do not assign the ID. It will be assigned by the session. id: str = '' """The unique identifier of the event.""" - timestamp: float = Field(default_factory=lambda: datetime.now().timestamp()) + timestamp: float = Field(default_factory=lambda: platform_time.get_time()) """The timestamp of the event.""" def model_post_init(self, __context): @@ -125,4 +127,4 @@ def has_trailing_code_execution_result( @staticmethod def new_id(): - return str(uuid.uuid4()) + return platform_uuid.new_uuid() diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 6f34e8fe75..8d6834e8c1 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -29,8 +29,8 @@ from typing import Dict from typing import Optional from typing import TYPE_CHECKING -import uuid +from google.adk.platform import uuid as platform_uuid from google.genai import types from ...agents.active_streaming_tool import ActiveStreamingTool @@ -175,7 +175,7 @@ def run_async_tool_in_new_loop(): def generate_client_function_call_id() -> str: - return f'{AF_FUNCTION_CALL_ID_PREFIX}{uuid.uuid4()}' + return f'{AF_FUNCTION_CALL_ID_PREFIX}{platform_uuid.new_uuid()}' def populate_client_function_call_id(model_response_event: Event) -> None: diff --git a/src/google/adk/platform/time.py b/src/google/adk/platform/time.py new file mode 100644 index 0000000000..e6b60aea90 --- /dev/null +++ b/src/google/adk/platform/time.py @@ -0,0 +1,43 @@ +# Copyright 2026 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. + +"""Platform module for abstracting system time generation.""" + +import time +from typing import Callable + +_default_time_provider: Callable[[], float] = time.time +_time_provider: Callable[[], float] = _default_time_provider + + +def set_time_provider(provider: Callable[[], float]) -> None: + """Sets the provider for the current time. + + Args: + provider: A callable that returns the current time in seconds since the + epoch. + """ + global _time_provider + _time_provider = provider + + +def reset_time_provider() -> None: + """Resets the time provider to its default implementation.""" + global _time_provider + _time_provider = _default_time_provider + + +def get_time() -> float: + """Returns the current time in seconds since the epoch.""" + return _time_provider() diff --git a/src/google/adk/platform/uuid.py b/src/google/adk/platform/uuid.py new file mode 100644 index 0000000000..ab6aedc6a0 --- /dev/null +++ b/src/google/adk/platform/uuid.py @@ -0,0 +1,42 @@ +# Copyright 2026 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. + +"""Platform module for abstracting unique ID generation.""" + +import uuid +from typing import Callable + +_default_id_provider: Callable[[], str] = lambda: str(uuid.uuid4()) +_id_provider: Callable[[], str] = _default_id_provider + + +def set_id_provider(provider: Callable[[], str]) -> None: + """Sets the provider for generating unique IDs. + + Args: + provider: A callable that returns a unique ID string. + """ + global _id_provider + _id_provider = provider + + +def reset_id_provider() -> None: + """Resets the ID provider to its default implementation.""" + global _id_provider + _id_provider = _default_id_provider + + +def new_uuid() -> str: + """Returns a new unique ID.""" + return _id_provider() diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index d782072d6a..c91a836fe3 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/src/google/adk/sessions/in_memory_session_service.py @@ -22,6 +22,8 @@ from typing_extensions import override +from google.adk.platform import time as platform_time +from google.adk.platform import uuid as platform_uuid from . import _session_util from ..errors.already_exists_error import AlreadyExistsError from ..events.event import Event @@ -108,14 +110,14 @@ def _create_session_impl( session_id = ( session_id.strip() if session_id and session_id.strip() - else str(uuid.uuid4()) + else platform_uuid.new_uuid() ) session = Session( app_name=app_name, user_id=user_id, id=session_id, state=session_state or {}, - last_update_time=time.time(), + last_update_time=platform_time.get_time(), ) if app_name not in self.sessions: diff --git a/src/google/adk/sessions/sqlite_session_service.py b/src/google/adk/sessions/sqlite_session_service.py index d23c8278cf..33f3f71920 100644 --- a/src/google/adk/sessions/sqlite_session_service.py +++ b/src/google/adk/sessions/sqlite_session_service.py @@ -29,6 +29,9 @@ import aiosqlite from typing_extensions import override +from google.adk.platform import time as platform_time +from google.adk.platform import uuid as platform_uuid + from . import _session_util from ..errors.already_exists_error import AlreadyExistsError from ..events.event import Event @@ -165,8 +168,8 @@ async def create_session( if session_id: session_id = session_id.strip() if not session_id: - session_id = str(uuid.uuid4()) - now = time.time() + session_id = platform_uuid.new_uuid() + now = platform_time.get_time() async with self._get_db_connection() as db: # Check if session_id already exists diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index ec74f8abe3..f9790e9257 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -418,9 +418,9 @@ async def test_list_artifact_versions_and_get_artifact_version( ] with patch( - "google.adk.artifacts.base_artifact_service.datetime" - ) as mock_datetime: - mock_datetime.now.return_value = FIXED_DATETIME + "google.adk.artifacts.base_artifact_service.platform_time" + ) as mock_platform_time: + mock_platform_time.get_time.return_value = FIXED_DATETIME.timestamp() for i in range(4): custom_metadata = {"key": "value" + str(i)} @@ -505,9 +505,9 @@ async def test_list_artifact_versions_with_user_prefix( ] with patch( - "google.adk.artifacts.base_artifact_service.datetime" - ) as mock_datetime: - mock_datetime.now.return_value = FIXED_DATETIME + "google.adk.artifacts.base_artifact_service.platform_time" + ) as mock_platform_time: + mock_platform_time.get_time.return_value = FIXED_DATETIME.timestamp() for i in range(4): custom_metadata = {"key": "value" + str(i)} diff --git a/tests/unittests/platform/test_time.py b/tests/unittests/platform/test_time.py new file mode 100644 index 0000000000..badff75f2a --- /dev/null +++ b/tests/unittests/platform/test_time.py @@ -0,0 +1,40 @@ +# Copyright 2026 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. + +"""Unit tests for the platform time module.""" + +import time +import unittest + +from google.adk.platform import time as platform_time + + +class TestTime(unittest.TestCase): + + def tearDown(self): + # Reset provider to default after each test + platform_time.reset_time_provider() + + def test_default_time_provider(self): + # Verify it returns a float that is close to now + now = time.time() + rt_time = platform_time.get_time() + self.assertIsInstance(rt_time, float) + self.assertAlmostEqual(rt_time, now, delta=1.0) + + def test_custom_time_provider(self): + # Test override + mock_time = 123456789.0 + platform_time.set_time_provider(lambda: mock_time) + self.assertEqual(platform_time.get_time(), mock_time) diff --git a/tests/unittests/platform/test_uuid.py b/tests/unittests/platform/test_uuid.py new file mode 100644 index 0000000000..afc8e378a7 --- /dev/null +++ b/tests/unittests/platform/test_uuid.py @@ -0,0 +1,40 @@ +# Copyright 2026 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. + +"""Unit tests for the platform uuid module.""" + +import uuid +import unittest + +from google.adk.platform import uuid as platform_uuid + + +class TestUUID(unittest.TestCase): + + def tearDown(self): + # Reset provider to default after each test + platform_uuid.reset_id_provider() + + def test_default_id_provider(self): + # Verify it returns a string uuid + uid = platform_uuid.new_uuid() + self.assertIsInstance(uid, str) + # Should be parseable as uuid + uuid.UUID(uid) + + def test_custom_id_provider(self): + # Test override + mock_id = "test-id-123" + platform_uuid.set_id_provider(lambda: mock_id) + self.assertEqual(platform_uuid.new_uuid(), mock_id)