gemini-code-assist[bot] commented on code in PR #35709: URL: https://github.com/apache/beam/pull/35709#discussion_r2236045934
########## sdks/python/apache_beam/examples/inference/gemini_text_classification.py: ########## @@ -67,8 +67,18 @@ def parse_known_args(argv): class PostProcessor(beam.DoFn): def process(self, element: PredictionResult) -> Iterable[str]: - yield "Input: " + str(element.example) + " Output: " + str( - element.inference[1][0].content.parts[0].text) + + inference = getattr(element, "inference", None) + + if hasattr(inference[1], "content"): + yield inference[1].content.parts[0].text + return + + if isinstance(inference[1], (tuple, list)) and len(inference) > 1: + yield "Input: " + str(element.example) + " Output: " + str( + inference[1][0].content.parts[0].text) + else: + yield "Can't decode inference for element: " + str(element.example) Review Comment:  The current implementation is not robust against different shapes of the `inference` object and could lead to runtime errors. 1. `inference` is retrieved using `getattr(element, "inference", None)`. If `inference` is `None`, any subsequent access like `inference[1]` on lines 73 and 77 will raise a `TypeError`. 2. If `inference` is not a sequence or has fewer than two elements, `inference[1]` will raise a `TypeError` or `IndexError`. 3. On line 77, `isinstance(inference[1], (tuple, list))` is checked, but it's not verified if the list is non-empty before accessing `inference[1][0]`. This could cause an `IndexError`. The condition `len(inference) > 1` does not protect against this, as it checks the length of the parent `inference` object. To prevent these potential crashes, it's best to validate the structure of `inference` before using it. The suggested change introduces these checks for safer execution. ```python inference = getattr(element, "inference", None) if not isinstance(inference, (tuple, list)) or len(inference) < 2: yield "Can't decode inference for element: " + str(element.example) return prediction_payload = inference[1] if hasattr(prediction_payload, "content"): yield prediction_payload.content.parts[0].text return if isinstance(prediction_payload, (tuple, list)) and prediction_payload: yield "Input: " + str(element.example) + " Output: " + str( prediction_payload[0].content.parts[0].text) else: yield "Can't decode inference for element: " + str(element.example) ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@beam.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
