From 9b76c722de505ae820024daa885e2f663529e774 Mon Sep 17 00:00:00 2001
From: Abram <israelvictory87@gmail.com>
Date: Tue, 28 May 2024 19:06:25 +0100
Subject: [PATCH] Refactor (instrument): improve decorator to handle async and
 sync funcs separately

---
 agenta-cli/agenta/sdk/decorators/tracing.py | 37 ++++++++++++++++-----
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/agenta-cli/agenta/sdk/decorators/tracing.py b/agenta-cli/agenta/sdk/decorators/tracing.py
index 4ea68bf692..8278ba4344 100644
--- a/agenta-cli/agenta/sdk/decorators/tracing.py
+++ b/agenta-cli/agenta/sdk/decorators/tracing.py
@@ -38,8 +38,10 @@ def __init__(
         self.tracing = ag.tracing
 
     def __call__(self, func: Callable[..., Any]):
+        is_coroutine_function = inspect.iscoroutinefunction(func)
+
         @wraps(func)
-        async def wrapper(*args, **kwargs):
+        async def async_wrapper(*args, **kwargs):
             result = None
             span = self.tracing.start_span(
                 name=func.__name__,
@@ -49,12 +51,7 @@ async def wrapper(*args, **kwargs):
             )
 
             try:
-                is_coroutine_function = inspect.iscoroutinefunction(func)
-                if is_coroutine_function:
-                    result = await func(*args, **kwargs)
-                else:
-                    result = func(*args, **kwargs)
-
+                result = await func(*args, **kwargs)
                 self.tracing.update_span_status(span=span, value="OK")
             except Exception as e:
                 result = str(e)
@@ -68,4 +65,28 @@ async def wrapper(*args, **kwargs):
                 )
             return result
 
-        return wrapper
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            result = None
+            span = self.tracing.start_span(
+                name=func.__name__,
+                input=kwargs,
+                spankind=self.spankind,
+                config=self.config,
+            )
+
+            try:
+                result = func(*args, **kwargs)
+                self.tracing.update_span_status(span=span, value="OK")
+            except Exception as e:
+                result = str(e)
+                self.tracing.update_span_status(span=span, value="ERROR")
+            finally:
+                self.tracing.end_span(
+                    outputs=(
+                        {"message": result} if not isinstance(result, dict) else result
+                    ),
+                    span=span,
+                )
+
+        return async_wrapper if is_coroutine_function else sync_wrapper