From 187ca85adaf74fcf226c0eb4fcb016ce5f37b2cf Mon Sep 17 00:00:00 2001
From: Nir Gazit <nirga@users.noreply.github.com>
Date: Mon, 20 May 2024 18:11:28 +0300
Subject: [PATCH] fix(sdk): serialize map outputs for non-promise outputs

---
 .../src/lib/tracing/decorators.ts             | 34 +++++++++++--------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
index 239bd1aa..c6e79bbc 100644
--- a/packages/traceloop-sdk/src/lib/tracing/decorators.ts
+++ b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
@@ -7,6 +7,7 @@ import {
 } from "@traceloop/ai-semantic-conventions";
 import { withAssociationProperties } from "./association";
 import { shouldSendTraces } from ".";
+import { Telemetry } from "../telemetry/telemetry";
 
 export type DecoratorConfig = {
   name: string;
@@ -82,8 +83,8 @@ function withEntity<
                 }),
               );
             }
-          } catch {
-            /* empty */
+          } catch (error) {
+            Telemetry.getInstance().logException(error);
           }
         }
 
@@ -92,19 +93,14 @@ function withEntity<
           return res.then((resolvedRes) => {
             try {
               if (shouldSendTraces()) {
-                if (resolvedRes instanceof Map) {
-                  span.setAttribute(
-                    SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
-                    JSON.stringify(Array.from(resolvedRes.entries())),
-                  );
-                } else {
-                  span.setAttribute(
-                    SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
-                    JSON.stringify(resolvedRes),
-                  );
-                }
+                span.setAttribute(
+                  SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
+                  serialize(resolvedRes),
+                );
               }
               return resolvedRes;
+            } catch (error) {
+              Telemetry.getInstance().logException(error);
             } finally {
               span.end();
             }
@@ -114,10 +110,12 @@ function withEntity<
           if (shouldSendTraces()) {
             span.setAttribute(
               SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
-              JSON.stringify(res),
+              serialize(res),
             );
           }
           return res;
+        } catch (error) {
+          Telemetry.getInstance().logException(error);
         } finally {
           span.end();
         }
@@ -240,3 +238,11 @@ export function tool(
 ) {
   return entity(TraceloopSpanKindValues.TOOL, config ?? {});
 }
+
+function serialize(input: unknown): string {
+  if (input instanceof Map) {
+    return JSON.stringify(Array.from(input.entries()));
+  } else {
+    return JSON.stringify(input);
+  }
+}