feat: Remove use of raw_event field in vertex ai session service

PiperOrigin-RevId: 893175221

Author: google-genai-bot

src/google/adk/sessions/vertex_ai_session_service.py

@@ -14,7 +14,6 @@
 from __future__ import annotations
 
 import asyncio
-import copy
 import datetime
 import json
 import logging
@@ -267,11 +266,72 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
     reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
 
-    raw_event_dict = event.model_dump(
-        exclude_none=True,
-        mode='json',
-        by_alias=True,
-    )
+    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
 
     async with self._get_api_client() as api_client:
       await api_client.agent_engines.sessions.events.append(
@@ -281,7 +341,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
           timestamp=datetime.datetime.fromtimestamp(
               event.timestamp, tz=datetime.timezone.utc
           ),
-          config={'raw_event': raw_event_dict},
+          config=config,
       )
     return event
 
@@ -330,22 +390,6 @@ 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:

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -28,7 +28,6 @@
 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
@@ -163,96 +162,6 @@ 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',
@@ -340,8 +249,6 @@ def _convert_to_object(data):
           'artifact_delta',
           'custom_metadata',
           'requested_auth_configs',
-          'rawEvent',
-          'raw_event',
       ]:
         kwargs[key] = value
       else:
@@ -773,38 +680,6 @@ 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):
@@ -941,36 +816,6 @@ 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)