feat: Use raw_event field in vertex ai session service for append and list events

google-genai-bot · copybara-github · commit 6ee036292e9e · 2026-03-31T13:56:23.000-07:00
PiperOrigin-RevId: 892504394
diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import datetime
 import json
 import logging
@@ -266,72 +267,11 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
     reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
 
-    config = {}
-    if event.content:
-      config['content'] = event.content.model_dump(
-          exclude_none=True, mode='json'
-      )
-    if event.actions:
-      config['actions'] = {
-          'skip_summarization': event.actions.skip_summarization,
-          'state_delta': event.actions.state_delta,
-          'artifact_delta': event.actions.artifact_delta,
-          'transfer_agent': event.actions.transfer_to_agent,
-          'escalate': event.actions.escalate,
-          'requested_auth_configs': {
-              k: json.loads(v.model_dump_json(exclude_none=True, by_alias=True))
-              for k, v in event.actions.requested_auth_configs.items()
-          },
-          # TODO: add requested_tool_confirmations, agent_state once
-          # they are available in the API.
-          # Note: compaction is stored via event_metadata.custom_metadata.
-      }
-    if event.error_code:
-      config['error_code'] = event.error_code
-    if event.error_message:
-      config['error_message'] = event.error_message
-
-    metadata_dict = {
-        'partial': event.partial,
-        'turn_complete': event.turn_complete,
-        'interrupted': event.interrupted,
-        'branch': event.branch,
-        'custom_metadata': event.custom_metadata,
-        'long_running_tool_ids': (
-            list(event.long_running_tool_ids)
-            if event.long_running_tool_ids
-            else None
-        ),
-    }
-    if event.grounding_metadata:
-      metadata_dict['grounding_metadata'] = event.grounding_metadata.model_dump(
-          exclude_none=True, mode='json'
-      )
-    # Store compaction data in custom_metadata since the Vertex AI service
-    # does not yet support the compaction field.
-    # TODO: Stop writing to custom_metadata once the Vertex AI service
-    # supports the compaction field natively in EventActions.
-    if event.actions and event.actions.compaction:
-      compaction_dict = event.actions.compaction.model_dump(
-          exclude_none=True, mode='json'
-      )
-      _set_internal_custom_metadata(
-          metadata_dict,
-          key=_COMPACTION_CUSTOM_METADATA_KEY,
-          value=compaction_dict,
-      )
-    # Store usage_metadata in custom_metadata since the Vertex AI service
-    # does not persist it in EventMetadata.
-    if event.usage_metadata:
-      usage_dict = event.usage_metadata.model_dump(
-          exclude_none=True, mode='json'
-      )
-      _set_internal_custom_metadata(
-          metadata_dict,
-          key=_USAGE_METADATA_CUSTOM_METADATA_KEY,
-          value=usage_dict,
-      )
-    config['event_metadata'] = metadata_dict
+    raw_event_dict = event.model_dump(
+        exclude_none=True,
+        mode='json',
+        by_alias=True,
+    )
 
     async with self._get_api_client() as api_client:
       await api_client.agent_engines.sessions.events.append(
@@ -341,7 +281,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
           timestamp=datetime.datetime.fromtimestamp(
               event.timestamp, tz=datetime.timezone.utc
           ),
-          config=config,
+          config={'raw_event': raw_event_dict},
       )
     return event
 
@@ -390,6 +330,22 @@ def _get_api_client(self) -> vertexai.AsyncClient:
 
 def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
   """Converts an API event object to an Event object."""
+  # Read event data from raw_event first before falling back to top level
+  # fields.
+  raw_event_dict = getattr(
+      api_event_obj, 'raw_event', getattr(api_event_obj, 'rawEvent', None)
+  )
+  if raw_event_dict:
+    event_dict = copy.deepcopy(raw_event_dict)
+    timestamp_obj = getattr(api_event_obj, 'timestamp', None)
+    event_dict.update({
+        'id': api_event_obj.name.split('/')[-1],
+        'invocation_id': getattr(api_event_obj, 'invocation_id', None),
+        'author': getattr(api_event_obj, 'author', None),
+        'timestamp': timestamp_obj.timestamp() if timestamp_obj else None,
+    })
+    return Event.model_validate(event_dict)
+
   actions = getattr(api_event_obj, 'actions', None)
   event_metadata = getattr(api_event_obj, 'event_metadata', None)
   if event_metadata:
diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py
@@ -28,6 +28,7 @@
 from google.adk.events.event import Event
 from google.adk.events.event_actions import EventActions
 from google.adk.events.event_actions import EventCompaction
+from google.adk.models.cache_metadata import CacheMetadata
 from google.adk.sessions.base_session_service import GetSessionConfig
 from google.adk.sessions.session import Session
 from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
@@ -162,6 +163,96 @@ def _generate_mock_events_for_session_5(num_events):
 MANY_EVENTS_COUNT = 200
 MOCK_EVENTS_JSON_5 = _generate_mock_events_for_session_5(MANY_EVENTS_COUNT)
 
+MOCK_EVENT_WITH_OVERRIDE_JSON = [{
+    'name': (
+        'projects/test-project/locations/test-location/'
+        'reasoningEngines/123/sessions/override/events/1'
+    ),
+    'invocationId': 'override_invoke',
+    'author': 'user_with_override',
+    'timestamp': '2024-12-12T12:12:12.123456Z',
+    'content': {
+        'parts': [
+            {'text': 'top_level_content'},
+        ],
+    },
+    'actions': {
+        'transferToAgent': 'top_level_agent',
+    },
+    'eventMetadata': {
+        'partial': True,
+        'turnComplete': False,
+        'interrupted': False,
+        'branch': 'top_level_branch',
+    },
+    'errorCode': '111',
+    'errorMessage': 'top_level_error',
+    'rawEvent': {
+        'invocationId': 'wrong_invocation_id',
+        'author': 'wrong_author',
+        'content': {
+            'parts': [
+                {'text': 'raw_event_content'},
+            ],
+        },
+        'actions': {
+            'transferToAgent': 'raw_event_agent',
+        },
+        'partial': False,
+        'turnComplete': True,
+        'interrupted': True,
+        'branch': 'raw_event_branch',
+        'errorCode': '222',
+        'errorMessage': 'raw_event_error',
+    },
+}]
+
+MOCK_EVENT_WITH_OVERRIDE_JSON_2 = [{
+    'name': (
+        'projects/test-project/locations/test-location/'
+        'reasoningEngines/123/sessions/override/events/1'
+    ),
+    'invocationId': 'override_invoke',
+    'author': 'user_with_override',
+    'content': {},
+    'actions': {},
+    'timestamp': '2024-12-12T12:12:12.123456Z',
+    'rawEvent': {
+        'invocationId': 'wrong_invocation_id',
+        'author': 'wrong_author',
+        'content': {
+            'parts': [
+                {'text': 'raw_event_content'},
+            ],
+        },
+        'actions': {
+            'skipSummarization': None,
+            'stateDelta': {},
+            'artifactDelta': {},
+            'transferToAgent': 'raw_event_agent',
+            'escalate': None,
+            'requestedAuthConfigs': {},
+        },
+        'errorCode': '222',
+        'errorMessage': 'raw_event_error',
+        'partial': False,
+        'turnComplete': True,
+        'interrupted': True,
+        'branch': 'raw_event_branch',
+        'customMetadata': None,
+        'longRunningToolIds': None,
+    },
+}]
+
+MOCK_SESSION_WITH_OVERRIDE_JSON = {
+    'name': (
+        'projects/test-project/locations/test-location/'
+        'reasoningEngines/123/sessions/override'
+    ),
+    'update_time': '2024-12-12T12:12:12.123456Z',
+    'user_id': 'user_with_override',
+}
+
 MOCK_SESSION = Session(
     app_name='123',
     user_id='user',
@@ -249,6 +340,8 @@ def _convert_to_object(data):
           'artifact_delta',
           'custom_metadata',
           'requested_auth_configs',
+          'rawEvent',
+          'raw_event',
       ]:
         kwargs[key] = value
       else:
@@ -680,6 +773,38 @@ async def test_get_session_keeps_events_newer_than_update_time(
   )
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+@pytest.mark.parametrize(
+    'mock_event_json',
+    [MOCK_EVENT_WITH_OVERRIDE_JSON, MOCK_EVENT_WITH_OVERRIDE_JSON_2],
+)
+async def test_get_session_from_raw_event(
+    mock_api_client_instance: MockAsyncClient,
+    mock_event_json,
+) -> None:
+  mock_api_client_instance.session_dict['6'] = MOCK_SESSION_WITH_OVERRIDE_JSON
+  mock_api_client_instance.event_dict['6'] = (
+      copy.deepcopy(mock_event_json),
+      None,
+  )
+  session_service = mock_vertex_ai_session_service()
+  session = await session_service.get_session(
+      app_name='123', user_id='user_with_override', session_id='6'
+  )
+  assert session is not None
+  assert len(session.events) == 1
+  event = session.events[0]
+  assert event.content.parts[0].text == 'raw_event_content'
+  assert event.actions.transfer_to_agent == 'raw_event_agent'
+  assert not event.partial
+  assert event.turn_complete
+  assert event.interrupted
+  assert event.branch == 'raw_event_branch'
+  assert event.error_code == '222'
+  assert event.error_message == 'raw_event_error'
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
 async def test_get_session_with_many_events(mock_api_client_instance):
@@ -816,6 +941,36 @@ async def test_append_event():
       branch='test_branch',
       custom_metadata={'custom': 'data'},
       long_running_tool_ids={'tool2'},
+      input_transcription=genai_types.Transcription(
+          text='test_input_transcription'
+      ),
+      output_transcription=genai_types.Transcription(
+          text='test_output_transcription'
+      ),
+      model_version='test_model_version',
+      avg_logprobs=0.5,
+      logprobs_result=genai_types.LogprobsResult(
+          chosen_candidates=[
+              genai_types.LogprobsResultCandidate(
+                  log_probability=0.5,
+                  token='test_token',
+                  token_id=0,
+              )
+          ]
+      ),
+      cache_metadata=CacheMetadata(
+          cache_name='test_cache_name',
+          fingerprint='test_fingerprint',
+          contents_count=1,
+      ),
+      citation_metadata=genai_types.CitationMetadata(
+          citations=[
+              genai_types.Citation(
+                  uri='http://test.com',
+                  title='test_title',
+              )
+          ]
+      ),
   )
 
   await session_service.append_event(session_before_append, event_to_append)
