From 639b51e862f581200bf49781b16a36f4fca6894a Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Thu, 12 Oct 2023 10:53:59 -0700
Subject: [PATCH] `lf.structured.ValueJsonRepr.parse` to pass through keyword
 args to `pg.from_json_str`.

This allows passing `force_dict=True` when dealing with hallucinated `_type` field.

PiperOrigin-RevId: 572954123
---
 langfun/core/structured/schema.py      | 2 +-
 langfun/core/structured/schema_test.py | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/langfun/core/structured/schema.py b/langfun/core/structured/schema.py
index 62444aac..76e09f95 100644
--- a/langfun/core/structured/schema.py
+++ b/langfun/core/structured/schema.py
@@ -605,7 +605,7 @@ def parse(self, text: str, schema: Schema | None = None, **kwargs) -> Any:
     del schema
     try:
       text = self.cleanup_json(text)
-      v = pg.from_json_str(text)
+      v = pg.from_json_str(text, **kwargs)
     except Exception as e:
       raise JsonError(text, e)  # pylint: disable=raise-missing-from
 
diff --git a/langfun/core/structured/schema_test.py b/langfun/core/structured/schema_test.py
index e6a94de3..eb42f9ef 100644
--- a/langfun/core/structured/schema_test.py
+++ b/langfun/core/structured/schema_test.py
@@ -162,6 +162,11 @@ def test_value_repr(self):
   def test_parse(self):
     schema = schema_lib.Schema(int)
     self.assertEqual(schema.parse('{"result": 1}'), 1)
+    schema = schema_lib.Schema(dict[str, int])
+    self.assertEqual(
+        schema.parse(
+            '{"result": {"_type": "Unknown", "x": 1}}}', force_dict=True),
+        dict(x=1))
     with self.assertRaisesRegex(
         schema_lib.SchemaError, 'Expect .* but encountered .*'):
       schema.parse('{"result": "def"}')