From 0bfc70be63521ed02e2bdd3393d945cbb9b86a46 Mon Sep 17 00:00:00 2001
From: Juan Pablo Vega <jp@agenta.ai>
Date: Fri, 5 Jul 2024 12:40:02 +0200
Subject: [PATCH] fix(sdk): AGE-272 Propagate func errors up in
 @ag.instrument() wrappers

Func errors used to be set as func result. Now they are propagated up by forwarding the exception up (raise e).
---
 .../agenta/sdk/decorators/llm_entrypoint.py   |  2 ++
 agenta-cli/agenta/sdk/decorators/tracing.py   | 33 ++++++++++---------
 agenta-cli/pyproject.toml                     |  2 +-
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/agenta-cli/agenta/sdk/decorators/llm_entrypoint.py b/agenta-cli/agenta/sdk/decorators/llm_entrypoint.py
index e18ae1bc41..4ee869222b 100644
--- 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
index a158740494..b7571477c8 100644
--- a/agenta-cli/agenta/sdk/decorators/tracing.py
+++ b/agenta-cli/agenta/sdk/decorators/tracing.py
@@ -58,19 +58,20 @@ 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:
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
+                self.tracing.update_span_status(span=span, value="ERROR")
+                self.tracing.end_span(outputs={})
+                raise e
 
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
@@ -89,17 +90,19 @@ 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:
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
+                self.tracing.update_span_status(span=span, value="ERROR")
+                self.tracing.end_span(outputs={})
+                raise e
 
         return async_wrapper if is_coroutine_function else sync_wrapper
diff --git a/agenta-cli/pyproject.toml b/agenta-cli/pyproject.toml
index 3ba8848ef2..966176c9c1 100644
--- a/agenta-cli/pyproject.toml
+++ b/agenta-cli/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.17.5"
+version = "0.17.6-alpha.1"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]