Merge branch 'main' into feat/parallelize-llm-judge-evaluation

Author: aryanpatel2121

src/google/adk/a2a/converters/event_converter.py

@@ -131,6 +131,7 @@ def _get_context_metadata(
         _get_adk_metadata_key("session_id"): invocation_context.session.id,
         _get_adk_metadata_key("invocation_id"): event.invocation_id,
         _get_adk_metadata_key("author"): event.author,
+        _get_adk_metadata_key("event_id"): event.id,
     }
 
     # Add optional metadata fields if present
@@ -479,7 +480,6 @@ def _create_status_update_event(
     task_id: Optional task ID to use for generated events.
     context_id: Optional Context ID to use for generated events.
 
-
   Returns:
     A TaskStatusUpdateEvent with RUNNING state.
   """

src/google/adk/evaluation/custom_metric_evaluator.py

@@ -24,7 +24,6 @@
 from .eval_case import ConversationScenario
 from .eval_case import Invocation
 from .eval_metrics import EvalMetric
-from .eval_metrics import EvalStatus
 from .evaluator import EvaluationResult
 from .evaluator import Evaluator
 
@@ -44,12 +43,6 @@ def _get_metric_function(
     ) from e
 
 
-def _get_eval_status(score: Optional[float], threshold: float) -> EvalStatus:
-  if score is None:
-    return EvalStatus.NOT_EVALUATED
-  return EvalStatus.PASSED if score >= threshold else EvalStatus.FAILED
-
-
 class _CustomMetricEvaluator(Evaluator):
   """Evaluator for custom metrics."""
 
@@ -64,16 +57,20 @@ async def evaluate_invocations(
       expected_invocations: Optional[list[Invocation]],
       conversation_scenario: Optional[ConversationScenario] = None,
   ) -> EvaluationResult:
+    eval_metric = self._eval_metric.model_copy(deep=True)
+    eval_metric.threshold = None
     if inspect.iscoroutinefunction(self._metric_function):
       eval_result = await self._metric_function(
-          actual_invocations, expected_invocations, conversation_scenario
+          eval_metric,
+          actual_invocations,
+          expected_invocations,
+          conversation_scenario,
       )
     else:
       eval_result = self._metric_function(
-          actual_invocations, expected_invocations, conversation_scenario
+          eval_metric,
+          actual_invocations,
+          expected_invocations,
+          conversation_scenario,
       )
-
-    eval_result.overall_eval_status = _get_eval_status(
-        eval_result.overall_score, self._eval_metric.threshold
-    )
     return eval_result

src/google/adk/evaluation/eval_metrics.py

@@ -258,9 +258,11 @@ class EvalMetric(EvalBaseModel):
       description="The name of the metric.",
   )
 
-  threshold: float = Field(
+  threshold: Optional[float] = Field(
+      default=None,
       description=(
-          "A threshold value. Each metric decides how to interpret this"
+          "This field will be deprecated soon. Please use `criterion` instead."
+          " A threshold value. Each metric decides how to interpret this"
           " threshold."
       ),
   )

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/a2a/converters/test_event_converter.py

@@ -56,6 +56,7 @@ def setup_method(self):
     self.mock_invocation_context.artifact_service = self.mock_artifact_service
 
     self.mock_event = Mock(spec=Event)
+    self.mock_event.id = None
     self.mock_event.invocation_id = "test-invocation-id"
     self.mock_event.author = "test-author"
     self.mock_event.branch = None
@@ -130,6 +131,7 @@ def test_get_context_metadata_success(self):
         f"{ADK_METADATA_KEY_PREFIX}session_id",
         f"{ADK_METADATA_KEY_PREFIX}invocation_id",
         f"{ADK_METADATA_KEY_PREFIX}author",
+        f"{ADK_METADATA_KEY_PREFIX}event_id",
     ]
 
     for key in expected_keys: