diff --git a/packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts b/packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts
index 02093c1e..fdf378c4 100644
--- a/packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts
+++ b/packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts
@@ -492,6 +492,27 @@ const transformTelemetryMetadata = (
       attributes[SpanAttributes.TRACELOOP_SPAN_KIND] =
         TraceloopSpanKindValues.AGENT;
       attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] = agentName;
+
+      const inputMessages = attributes[SpanAttributes.LLM_INPUT_MESSAGES];
+      const outputMessages = attributes[SpanAttributes.LLM_OUTPUT_MESSAGES];
+      const toolArgs = attributes["ai.toolCall.args"];
+      const toolResult = attributes["ai.toolCall.result"];
+
+      if (inputMessages || outputMessages) {
+        if (inputMessages) {
+          attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT] = inputMessages;
+        }
+        if (outputMessages) {
+          attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = outputMessages;
+        }
+      } else if (toolArgs || toolResult) {
+        if (toolArgs) {
+          attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT] = toolArgs;
+        }
+        if (toolResult) {
+          attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = toolResult;
+        }
+      }
     }
   }