Fix VertexAI unittest failure due to SDK changes.

PiperOrigin-RevId: 696637774
google · Nov 14, 2024 · 5fc599e · 5fc599e
1 parent 5d5be83
commit 5fc599e
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/langfun/core/llms/vertexai_test.py b/langfun/core/llms/vertexai_test.py
@@ -199,6 +199,15 @@ def test_generation_config(self):
     # There is a discrepancy between the `property_ordering` in the
     # Google-internal version and the open-source version.
     actual['response_schema'].pop('property_ordering', None)
+    if pg.KeyPath('response_schema.type_').get(actual):
+      actual['response_schema']['type'] = actual['response_schema'].pop('type_')
+    if 'type_' in actual['response_schema']:
+      actual['response_schema']['type'] = actual['response_schema'].pop('type_')
+    if 'type_' in actual['response_schema']['properties']['name']:
+      actual['response_schema']['properties']['name'][
+          'type'
+      ] = actual['response_schema']['properties']['name'].pop('type_')
+
     self.assertEqual(
         actual,
         dict(
@@ -209,9 +218,9 @@ def test_generation_config(self):
             stop_sequences=['\n'],
             response_mime_type='application/json',
             response_schema={
-                'type_': 'OBJECT',
+                'type': 'OBJECT',
                 'properties': {
-                    'name': {'type_': 'STRING'}
+                    'name': {'type': 'STRING'}
                 },
                 'required': ['name'],
                 'title': 'Person',