From ab4efd880fd2e7a81171f8678becd13b2b625675 Mon Sep 17 00:00:00 2001 From: Daiyi Peng Date: Thu, 1 May 2025 15:01:03 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 753741877 --- langfun/core/structured/querying.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/langfun/core/structured/querying.py b/langfun/core/structured/querying.py index 646d5326..2bd05f8f 100644 --- a/langfun/core/structured/querying.py +++ b/langfun/core/structured/querying.py @@ -458,7 +458,9 @@ class Entity(pg.Object): if protocol is None: protocol = lf.context_value('__query_protocol__', 'python') - invocation_id = invocation_id or f'query@{uuid.uuid4().hex[-7:]}' + def _invocation_id(): + return invocation_id or f'query@{uuid.uuid4().hex[-7:]}' + # Multiple quries will be issued when `lm` is a list or `num_samples` is # greater than 1. if isinstance(lm, list) or num_samples != 1: @@ -482,7 +484,7 @@ def _single_query(inputs): protocol=protocol, returns_message=returns_message, skip_lm=skip_lm, - invocation_id=f'{invocation_id}:{i}', + invocation_id=f'{_invocation_id()}:{i}', **kwargs, ) lm_list = lm if isinstance(lm, list) else [lm] @@ -595,7 +597,7 @@ def _result(message: lf.Message): metadata.pop('usage', None) invocation = QueryInvocation( - id=invocation_id, + id=_invocation_id(), input=pg.Ref(query_input), schema=( schema_lib.Schema.from_value(schema)