Merge pull request #1846 from Agenta-AI/fix/llm-app-error-as-output-m…

…essage fix(sdk): AGE-272 Propagate func errors up in @ag.instrument() wrappers
Agenta-AI · Jul 5, 2024 · 3818fbf · 3818fbf
2 parents 2560667 + 2d2ee83
commit 3818fbf
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 15 deletions.
diff --git a/agenta-cli/agenta/sdk/decorators/llm_entrypoint.py b/agenta-cli/agenta/sdk/decorators/llm_entrypoint.py
@@ -83,9 +83,11 @@ async def wrapper(*args, **kwargs) -> Any:
                 {"config": config_params, "environment": "playground"}
             )
 
+            # Exceptions are all handled inside self.execute_function()
             llm_result = await self.execute_function(
                 func, *args, params=func_params, config_params=config_params
             )
+
             return llm_result
 
         @functools.wraps(func)

diff --git a/agenta-cli/agenta/sdk/decorators/tracing.py b/agenta-cli/agenta/sdk/decorators/tracing.py
@@ -58,19 +58,24 @@ async def async_wrapper(*args, **kwargs):
             try:
                 result = await func(*args, **kwargs)
                 self.tracing.update_span_status(span=span, value="OK")
-            except Exception as e:
-                result = str(e)
-                self.tracing.set_span_attribute(
-                    {"traceback_exception": traceback.format_exc()}
-                )
-                self.tracing.update_span_status(span=span, value="ERROR")
-            finally:
                 self.tracing.end_span(
                     outputs=(
                         {"message": result} if not isinstance(result, dict) else result
                     )
                 )
-            return result
+                return result
+
+            except Exception as e:
+                result = {
+                    "message": str(e),
+                    "stacktrace": traceback.format_exc(),
+                }
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
+                self.tracing.update_span_status(span=span, value="ERROR")
+                self.tracing.end_span(outputs=result)
+                raise e
 
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
@@ -89,17 +94,23 @@ def sync_wrapper(*args, **kwargs):
             try:
                 result = func(*args, **kwargs)
                 self.tracing.update_span_status(span=span, value="OK")
-            except Exception as e:
-                result = str(e)
-                self.tracing.set_span_attribute(
-                    {"traceback_exception": traceback.format_exc()}
-                )
-                self.tracing.update_span_status(span=span, value="ERROR")
-            finally:
                 self.tracing.end_span(
                     outputs=(
                         {"message": result} if not isinstance(result, dict) else result
                     )
                 )
+                return result
+
+            except Exception as e:
+                result = {
+                    "message": str(e),
+                    "stacktrace": traceback.format_exc(),
+                }
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
+                self.tracing.update_span_status(span=span, value="ERROR")
+                self.tracing.end_span(outputs=result)
+                raise e
 
         return async_wrapper if is_coroutine_function else sync_wrapper