diff --git a/langfun/core/structured/querying.py b/langfun/core/structured/querying.py index f7f0af0..d73c226 100644 --- a/langfun/core/structured/querying.py +++ b/langfun/core/structured/querying.py @@ -583,7 +583,16 @@ def lm_request(self) -> lf.Message: @functools.cached_property def output(self) -> Any: - return query_output(self.lm_response, self.schema) + """The output of `lf.query`. If it failed, returns the `MappingError`.""" + try: + return query_output(self.lm_response, self.schema) + except mapping.MappingError as e: + return e + + @property + def has_error(self) -> bool: + """Returns True if the query failed to generate a valid output.""" + return isinstance(self.output, BaseException) @property def elapse(self) -> float: diff --git a/langfun/core/structured/querying_test.py b/langfun/core/structured/querying_test.py index 8d32859..735ee41 100644 --- a/langfun/core/structured/querying_test.py +++ b/langfun/core/structured/querying_test.py @@ -1051,6 +1051,16 @@ def test_query(self): class QueryInvocationTest(unittest.TestCase): + def test_basics(self): + lm = fake.StaticSequence([ + 'Activity(description="hi"', + ]) + with querying.track_queries() as queries: + querying.query('foo', Activity, default=None, lm=lm) + + self.assertTrue(queries[0].has_error) + self.assertIsInstance(queries[0].output, mapping.MappingError) + def test_to_html(self): lm = fake.StaticSequence([ 'Activity(description="hi")',