fix(llamaindex): report inputs/outputs (#132)

Author: nirga

packages/instrumentation-llamaindex/recordings/Test-LlamaIndex-instrumentation_1988279490/should-add-span-for-all-instrumented-methods_2883459899/recording.har

@@ -23,7 +23,7 @@ import {
 
 import { LlamaIndexInstrumentationConfig } from "./types";
 import { CustomLLMInstrumentation } from "./custom-llm-instrumentation";
-import { genericWrapper } from "./utils";
+import { genericWrapper, shouldSendPrompts } from "./utils";
 
 import { BaseEmbedding, BaseSynthesizer, LLM, BaseRetriever } from "llamaindex";
 import { TraceloopSpanKindValues } from "@traceloop/ai-semantic-conventions";
@@ -102,6 +102,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
         "query",
         TraceloopSpanKindValues.WORKFLOW,
         this.tracer,
+        shouldSendPrompts(this._config),
       ),
     );
 
@@ -113,6 +114,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
         "chat",
         TraceloopSpanKindValues.WORKFLOW,
         this.tracer,
+        shouldSendPrompts(this._config),
       ),
     );
 
@@ -133,6 +135,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             "getQueryEmbedding",
             TraceloopSpanKindValues.TASK,
             this.tracer,
+            shouldSendPrompts(this._config),
           ),
         );
       } else if (this.isSynthesizer(cls.prototype)) {
@@ -144,6 +147,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             "synthesize",
             TraceloopSpanKindValues.TASK,
             this.tracer,
+            shouldSendPrompts(this._config),
           ),
         );
       } else if (this.isRetriever(cls.prototype)) {
@@ -155,6 +159,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             "retrieve",
             TraceloopSpanKindValues.TASK,
             this.tracer,
+            shouldSendPrompts(this._config),
           ),
         );
       }

packages/instrumentation-llamaindex/src/instrumentation.ts

@@ -15,7 +15,7 @@ export const shouldSendPrompts = (config: LlamaIndexInstrumentationConfig) => {
     .getValue(CONTEXT_KEY_ALLOW_TRACE_CONTENT);
 
   if (contextShouldSendPrompts !== undefined) {
-    return contextShouldSendPrompts;
+    return !!contextShouldSendPrompts;
   }
 
   return config.traceContent !== undefined ? config.traceContent : true;
@@ -45,6 +45,7 @@ export function genericWrapper(
   methodName: string,
   kind: TraceloopSpanKindValues,
   tracer: Tracer,
+  shouldSendPrompts: boolean,
 ) {
   // eslint-disable-next-line @typescript-eslint/ban-types
   return (original: Function) => {
@@ -57,6 +58,33 @@ export function genericWrapper(
         span.setAttribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, name);
       }
 
+      if (shouldSendPrompts) {
+        try {
+          if (
+            args.length === 1 &&
+            typeof args[0] === "object" &&
+            !(args[0] instanceof Map)
+          ) {
+            span.setAttribute(
+              SpanAttributes.TRACELOOP_ENTITY_INPUT,
+              JSON.stringify({ args: [], kwargs: args[0] }),
+            );
+          } else {
+            span.setAttribute(
+              SpanAttributes.TRACELOOP_ENTITY_INPUT,
+              JSON.stringify({
+                args: args.map((arg) =>
+                  arg instanceof Map ? Array.from(arg.entries()) : arg,
+                ),
+                kwargs: {},
+              }),
+            );
+          }
+        } catch {
+          /* empty */
+        }
+      }
+
       const execContext = trace.setSpan(context.active(), span);
       const execPromise = safeExecuteInTheMiddle(
         () => {
@@ -71,8 +99,25 @@ export function genericWrapper(
         .then((result: any) => {
           return new Promise((resolve) => {
             span.setStatus({ code: SpanStatusCode.OK });
-            span.end();
-            resolve(result);
+
+            try {
+              if (shouldSendPrompts) {
+                if (result instanceof Map) {
+                  span.setAttribute(
+                    SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
+                    JSON.stringify(Array.from(result.entries())),
+                  );
+                } else {
+                  span.setAttribute(
+                    SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
+                    JSON.stringify(result),
+                  );
+                }
+              }
+            } finally {
+              span.end();
+              resolve(result);
+            }
           });
         })
         .catch((error: Error) => {

packages/instrumentation-llamaindex/src/utils.ts

@@ -171,10 +171,33 @@ describe("Test LlamaIndex instrumentation", async function () {
     // TODO: Need to figure out why this doesn't get logged
     // assert.ok(spanNames.includes("get_query_embedding.task"));
 
-    assert.ok(spanNames.includes("retriever_query_engine.query"));
+    const retrieverQueryEngineSpan = spans.find(
+      (span) => span.name === "retriever_query_engine.query",
+    );
+
     assert.ok(spanNames.includes("retriever_query_engine.retrieve"));
     assert.ok(spanNames.includes("llamaindex.open_ai.chat"));
     assert.ok(spanNames.includes("response_synthesizer.synthesize"));
     assert.ok(spanNames.includes("vector_index_retriever.retrieve"));
+
+    assert.ok(retrieverQueryEngineSpan);
+    assert.ok(retrieverQueryEngineSpan.attributes["traceloop.entity.input"]);
+    assert.ok(retrieverQueryEngineSpan.attributes["traceloop.entity.output"]);
+    assert.strictEqual(
+      JSON.parse(
+        retrieverQueryEngineSpan.attributes[
+          "traceloop.entity.input"
+        ].toString(),
+      ).kwargs.query,
+      "Where was albert einstein born?",
+    );
+    assert.strictEqual(
+      JSON.parse(
+        retrieverQueryEngineSpan.attributes[
+          "traceloop.entity.output"
+        ].toString(),
+      ).response,
+      result.response,
+    );
   }).timeout(60000);
 });