Make stacktraces more readable by dropping excessive helper methods from them.

PiperOrigin-RevId: 809748640

Author: epiqueras

genai_processors/context.py

@@ -17,7 +17,6 @@
 import asyncio
 from collections.abc import Coroutine, Iterable
 import contextvars
-import types
 from typing import Any, TypeVar
 
 from absl import logging
@@ -39,52 +38,11 @@
 
 
 def raise_flattened_exception_group(exception: Exception):
-  """Flattens the exception group and filters the traceback."""
-  exceptions_in_group = 0
   e = exception
   while isinstance(e, ExceptionGroup):
-    exceptions_in_group += len(e.exceptions)
     e = e.exceptions[0]
-
-  # Every processor invocation is wrapperd in a couple of helpers, significantly
-  # inflating the callstack. Here we manipulate it and leave only the most
-  # nested instances of these helper, making stacktraces more readable.
-  processors_internal_functions = [
-      '_normalize_part_stream',
-      '_enqueue_content',
-      '__call__',
-  ]
-  observed_functions = set()
-
-  if e.__traceback__ is not None:
-    frames = []
-    tb = e.__traceback__
-    while tb is not None:
-      frames.append(tb)
-      tb = tb.tb_next
-
-    tb = None
-    for frame in reversed(frames):
-      if frame.tb_frame.f_code.co_filename.endswith(
-          'genai_processors/processor.py'
-      ):
-        function_name = frame.tb_frame.f_code.co_name
-        if function_name in processors_internal_functions:
-          if function_name in observed_functions:
-            continue
-          observed_functions.add(function_name)
-
-        tb = types.TracebackType(
-            tb_next=tb,
-            tb_frame=frame.tb_frame,
-            tb_lasti=frame.tb_lasti,
-            tb_lineno=frame.tb_lineno,
-        )
-
-    e.__traceback__ = tb
-
-  if e is exception or exceptions_in_group <= 1:
-    raise e from None
+  if e is exception:
+    raise exception
   else:
     raise e from exception
 

genai_processors/tests/processor_test.py

@@ -149,15 +149,18 @@ async def processor_0(
     # Using assertRaises removes the traceback which we want to check
     # includes the TaskGroup error.
     exception = None
+    tb = None
     try:
       processor.apply_sync(
           processor_0, content_api.ProcessorContent(['foo', 'bar'])
       )
     except ValueError as e:
       exception = str(e)
+      tb = traceback.format_exc()
 
     # Assert that the error message isn't a generic ExceptionGroup message
     self.assertNotIn('unhandled errors in a TaskGroup', exception)
+    self.assertIn('unhandled errors in a TaskGroup', tb)
     self.assertEqual(exception, err_msg)
 
   def test_chained_part_processor_raises_specific_error_message(self):
@@ -176,15 +179,18 @@ async def processor_0(
     # Using assertRaises removes the traceback which we want to check
     # includes the TaskGroup error.
     exception = None
+    tb = None
     try:
       processor.apply_sync(
           combined_processor, [content_api.ProcessorPart('bar')]
       )
     except ValueError as e:
       exception = str(e)
+      tb = traceback.format_exc()
 
     # Assert that the error message isn't a generic ExceptionGroup message
     self.assertNotIn('unhandled errors in a TaskGroup', exception)
+    self.assertIn('unhandled errors in a TaskGroup', tb)
     self.assertEqual(exception, err_msg)
 
   def test_normalization_function(self):
@@ -999,24 +1005,6 @@ async def proc_c(
     content = processor.apply_sync(chain_abc, inputs)
     self.assertEqual(content, inputs)
 
-  def test_backtrace_optimization(self):
-    @processor.part_processor_function
-    async def fail(
-        part: content_api.ProcessorPart,
-    ) -> AsyncIterable[content_api.ProcessorPart]:
-      raise ValueError('foo is not allowed')
-      yield part  # pylint: disable=unreachable
-
-    p = processor.passthrough()
-
-    inputs = [content_api.ProcessorPart('0')]
-    try:
-      processor.apply_sync(p + p + fail, inputs)
-    except ValueError:
-      tb = traceback.format_exc()
-      self.assertIn('foo is not allowed', tb)  # pylint: disable=g-assert-in-except
-      self.assertEqual(tb.count('_normalize_part_stream\n'), 1)  # pylint: disable=g-assert-in-except
-
 
 class ParallelProcessorsTest(TestWithProcessors, parameterized.TestCase):