Merge branch 'main' into mypy-gh-actions

Author: sararob

src/google/adk/evaluation/vertex_ai_eval_facade.py

@@ -159,18 +159,25 @@ def _perform_eval(dataset, metrics):
     """
     project_id = os.environ.get("GOOGLE_CLOUD_PROJECT", None)
     location = os.environ.get("GOOGLE_CLOUD_LOCATION", None)
-
-    if not project_id:
-      raise ValueError("Missing project id." + _ERROR_MESSAGE_SUFFIX)
-    if not location:
-      raise ValueError("Missing location." + _ERROR_MESSAGE_SUFFIX)
-
-    from vertexai import Client
-    from vertexai import types as vertexai_types
-
-    client = Client(project=project_id, location=location)
+    api_key = os.environ.get("GOOGLE_API_KEY", None)
+
+    from ..dependencies.vertexai import vertexai
+
+    if api_key:
+      client = vertexai.Client(api_key=api_key)
+    elif project_id or location:
+      if not project_id:
+        raise ValueError("Missing project id." + _ERROR_MESSAGE_SUFFIX)
+      if not location:
+        raise ValueError("Missing location." + _ERROR_MESSAGE_SUFFIX)
+      client = vertexai.Client(project=project_id, location=location)
+    else:
+      raise ValueError(
+          "Either API Key or Google cloud Project id and location should be"
+          " specified."
+      )
 
     return client.evals.evaluate(
-        dataset=vertexai_types.EvaluationDataset(eval_dataset_df=dataset),
+        dataset=vertexai.types.EvaluationDataset(eval_dataset_df=dataset),
         metrics=metrics,
     )

src/google/adk/sessions/database_session_service.py

@@ -31,7 +31,6 @@
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
 from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.inspection import inspect
 from sqlalchemy.pool import StaticPool
 from typing_extensions import override
 from tzlocal import get_localzone
@@ -275,22 +274,29 @@ async def create_session(
         storage_user_state.state = storage_user_state.state | user_state_delta
 
       # Store the session
+      now = datetime.now(timezone.utc)
+      is_sqlite = self.db_engine.dialect.name == "sqlite"
+      if is_sqlite:
+        now = now.replace(tzinfo=None)
+
       storage_session = schema.StorageSession(
           app_name=app_name,
           user_id=user_id,
           id=session_id,
           state=session_state,
+          create_time=now,
+          update_time=now,
       )
       sql_session.add(storage_session)
       await sql_session.commit()
 
-      await sql_session.refresh(storage_session)
-
       # Merge states for response
       merged_state = _merge_state(
           storage_app_state.state, storage_user_state.state, session_state
       )
-      session = storage_session.to_session(state=merged_state)
+      session = storage_session.to_session(
+          state=merged_state, is_sqlite=is_sqlite
+      )
     return session
 
   @override
@@ -350,7 +356,10 @@ async def get_session(
 
       # Convert storage session to session
       events = [e.to_event() for e in reversed(storage_events)]
-      session = storage_session.to_session(state=merged_state, events=events)
+      is_sqlite = self.db_engine.dialect.name == "sqlite"
+      session = storage_session.to_session(
+          state=merged_state, events=events, is_sqlite=is_sqlite
+      )
     return session
 
   @override
@@ -393,11 +402,14 @@ async def list_sessions(
           user_states_map[storage_user_state.user_id] = storage_user_state.state
 
       sessions = []
+      is_sqlite = self.db_engine.dialect.name == "sqlite"
       for storage_session in results:
         session_state = storage_session.state
         user_state = user_states_map.get(storage_session.user_id, {})
         merged_state = _merge_state(app_state, user_state, session_state)
-        sessions.append(storage_session.to_session(state=merged_state))
+        sessions.append(
+            storage_session.to_session(state=merged_state, is_sqlite=is_sqlite)
+        )
       return ListSessionsResponse(sessions=sessions)
 
   @override
@@ -433,15 +445,6 @@ async def append_event(self, session: Session, event: Event) -> Event:
           schema.StorageSession, (session.app_name, session.user_id, session.id)
       )
 
-      if storage_session.update_timestamp_tz > session.last_update_time:
-        raise ValueError(
-            "The last_update_time provided in the session object"
-            f" {datetime.fromtimestamp(session.last_update_time):'%Y-%m-%d %H:%M:%S'} is"
-            " earlier than the update_time in the storage_session"
-            f" {datetime.fromtimestamp(storage_session.update_timestamp_tz):'%Y-%m-%d %H:%M:%S'}."
-            " Please check if it is a stale session."
-        )
-
       # Fetch states from storage
       storage_app_state = await sql_session.get(
           schema.StorageAppState, (session.app_name)
@@ -450,6 +453,29 @@ async def append_event(self, session: Session, event: Event) -> Event:
           schema.StorageUserState, (session.app_name, session.user_id)
       )
 
+      is_sqlite = self.db_engine.dialect.name == "sqlite"
+      if (
+          storage_session.get_update_timestamp(is_sqlite)
+          > session.last_update_time
+      ):
+        # Reload the session from storage if it has been updated since it was
+        # loaded.
+        app_state = storage_app_state.state if storage_app_state else {}
+        user_state = storage_user_state.state if storage_user_state else {}
+        session_state = storage_session.state
+        session.state = _merge_state(app_state, user_state, session_state)
+
+        stmt = (
+            select(schema.StorageEvent)
+            .filter(schema.StorageEvent.app_name == session.app_name)
+            .filter(schema.StorageEvent.session_id == session.id)
+            .filter(schema.StorageEvent.user_id == session.user_id)
+            .order_by(schema.StorageEvent.timestamp.asc())
+        )
+        result = await sql_session.stream_scalars(stmt)
+        storage_events = [e async for e in result]
+        session.events = [e.to_event() for e in storage_events]
+
       # Extract state delta
       if event.actions and event.actions.state_delta:
         state_deltas = _session_util.extract_state_delta(
@@ -466,7 +492,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
         if session_state_delta:
           storage_session.state = storage_session.state | session_state_delta
 
-      if storage_session._dialect_name == "sqlite":
+      if is_sqlite:
         update_time = datetime.fromtimestamp(
             event.timestamp, timezone.utc
         ).replace(tzinfo=None)
@@ -476,10 +502,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
       sql_session.add(schema.StorageEvent.from_event(session, event))
 
       await sql_session.commit()
-      await sql_session.refresh(storage_session)
 
       # Update timestamp with commit time
-      session.last_update_time = storage_session.update_timestamp_tz
+      session.last_update_time = storage_session.get_update_timestamp(is_sqlite)
 
     # Also update the in-memory session
     await super().append_event(session=session, event=event)

src/google/adk/sessions/schemas/v0.py

@@ -41,7 +41,6 @@
 from sqlalchemy import Text
 from sqlalchemy.dialects import mysql
 from sqlalchemy.ext.mutable import MutableDict
-from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
@@ -131,15 +130,9 @@ class StorageSession(Base):
   def __repr__(self):
     return f"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @property
-  def _dialect_name(self) -> Optional[str]:
-    session = inspect(self).session
-    return session.bind.dialect.name if session else None
-
-  @property
-  def update_timestamp_tz(self) -> datetime:
+  def get_update_timestamp(self, is_sqlite: bool) -> float:
     """Returns the time zone aware update timestamp."""
-    if self._dialect_name == "sqlite":
+    if is_sqlite:
       # SQLite does not support timezone. SQLAlchemy returns a naive datetime
       # object without timezone information. We need to convert it to UTC
       # manually.
@@ -150,6 +143,7 @@ def to_session(
       self,
       state: dict[str, Any] | None = None,
       events: list[Event] | None = None,
+      is_sqlite: bool = False,
   ) -> Session:
     """Converts the storage session to a session object."""
     if state is None:
@@ -163,7 +157,7 @@ def to_session(
         id=self.id,
         state=state,
         events=events,
-        last_update_time=self.update_timestamp_tz,
+        last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
     )
 
 

src/google/adk/sessions/schemas/v1.py

@@ -32,7 +32,6 @@
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import func
 from sqlalchemy.ext.mutable import MutableDict
-from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
@@ -106,15 +105,9 @@ class StorageSession(Base):
   def __repr__(self):
     return f"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
-  @property
-  def _dialect_name(self) -> Optional[str]:
-    session = inspect(self).session
-    return session.bind.dialect.name if session else None
-
-  @property
-  def update_timestamp_tz(self) -> datetime:
+  def get_update_timestamp(self, is_sqlite: bool) -> float:
     """Returns the time zone aware update timestamp."""
-    if self._dialect_name == "sqlite":
+    if is_sqlite:
       # SQLite does not support timezone. SQLAlchemy returns a naive datetime
       # object without timezone information. We need to convert it to UTC
       # manually.
@@ -125,6 +118,7 @@ def to_session(
       self,
       state: dict[str, Any] | None = None,
       events: list[Event] | None = None,
+      is_sqlite: bool = False,
   ) -> Session:
     """Converts the storage session to a session object."""
     if state is None:
@@ -138,7 +132,7 @@ def to_session(
         id=self.id,
         state=state,
         events=events,
-        last_update_time=self.update_timestamp_tz,
+        last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
     )
 
 

tests/unittests/evaluation/test_vertex_ai_eval_facade.py

@@ -16,13 +16,15 @@
 
 """Tests for the Response Evaluator."""
 import math
+import os
 import random
 
 from google.adk.dependencies.vertexai import vertexai
 from google.adk.evaluation.eval_case import Invocation
 from google.adk.evaluation.evaluator import EvalStatus
 from google.adk.evaluation.vertex_ai_eval_facade import _VertexAiEvalFacade
 from google.genai import types as genai_types
+import pandas as pd
 import pytest
 
 vertexai_types = vertexai.types
@@ -246,3 +248,89 @@ def test_evaluate_invocations_metric_multiple_invocations(self, mocker):
     )
     assert evaluation_result.overall_eval_status == EvalStatus.FAILED
     assert mock_perform_eval.call_count == num_invocations
+
+  def test_perform_eval_with_api_key(self, mocker):
+    mocker.patch.dict(
+        os.environ, {"GOOGLE_API_KEY": "test_api_key"}, clear=True
+    )
+    mock_client_cls = mocker.patch(
+        "google.adk.dependencies.vertexai.vertexai.Client"
+    )
+    mock_client_instance = mock_client_cls.return_value
+    dummy_dataset = pd.DataFrame(
+        [{"prompt": "p", "reference": "r", "response": "r"}]
+    )
+    dummy_metrics = [vertexai_types.PrebuiltMetric.COHERENCE]
+
+    _VertexAiEvalFacade._perform_eval(dummy_dataset, dummy_metrics)
+
+    mock_client_cls.assert_called_once_with(api_key="test_api_key")
+    mock_client_instance.evals.evaluate.assert_called_once()
+
+  def test_perform_eval_with_project_and_location(self, mocker):
+    mocker.patch.dict(
+        os.environ,
+        {
+            "GOOGLE_CLOUD_PROJECT": "test_project",
+            "GOOGLE_CLOUD_LOCATION": "test_location",
+        },
+        clear=True,
+    )
+    mock_client_cls = mocker.patch(
+        "google.adk.dependencies.vertexai.vertexai.Client"
+    )
+    mock_client_instance = mock_client_cls.return_value
+    dummy_dataset = pd.DataFrame(
+        [{"prompt": "p", "reference": "r", "response": "r"}]
+    )
+    dummy_metrics = [vertexai_types.PrebuiltMetric.COHERENCE]
+
+    _VertexAiEvalFacade._perform_eval(dummy_dataset, dummy_metrics)
+
+    mock_client_cls.assert_called_once_with(
+        project="test_project", location="test_location"
+    )
+    mock_client_instance.evals.evaluate.assert_called_once()
+
+  def test_perform_eval_with_project_only_raises_error(self, mocker):
+    mocker.patch.dict(
+        os.environ, {"GOOGLE_CLOUD_PROJECT": "test_project"}, clear=True
+    )
+    mocker.patch("google.adk.dependencies.vertexai.vertexai.Client")
+    dummy_dataset = pd.DataFrame(
+        [{"prompt": "p", "reference": "r", "response": "r"}]
+    )
+    dummy_metrics = [vertexai_types.PrebuiltMetric.COHERENCE]
+
+    with pytest.raises(ValueError, match="Missing location."):
+      _VertexAiEvalFacade._perform_eval(dummy_dataset, dummy_metrics)
+
+  def test_perform_eval_with_location_only_raises_error(self, mocker):
+    mocker.patch.dict(
+        os.environ, {"GOOGLE_CLOUD_LOCATION": "test_location"}, clear=True
+    )
+    mocker.patch("google.adk.dependencies.vertexai.vertexai.Client")
+    dummy_dataset = pd.DataFrame(
+        [{"prompt": "p", "reference": "r", "response": "r"}]
+    )
+    dummy_metrics = [vertexai_types.PrebuiltMetric.COHERENCE]
+
+    with pytest.raises(ValueError, match="Missing project id."):
+      _VertexAiEvalFacade._perform_eval(dummy_dataset, dummy_metrics)
+
+  def test_perform_eval_with_no_env_vars_raises_error(self, mocker):
+    mocker.patch.dict(os.environ, {}, clear=True)
+    mocker.patch("google.adk.dependencies.vertexai.vertexai.Client")
+    dummy_dataset = pd.DataFrame(
+        [{"prompt": "p", "reference": "r", "response": "r"}]
+    )
+    dummy_metrics = [vertexai_types.PrebuiltMetric.COHERENCE]
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            "Either API Key or Google cloud Project id and location should be"
+            " specified."
+        ),
+    ):
+      _VertexAiEvalFacade._perform_eval(dummy_dataset, dummy_metrics)