From 586e66fa88c77bcb2e4a8bfe313d02ba2b31039e Mon Sep 17 00:00:00 2001
From: Milan Gallas <44726162+GALLLASMILAN@users.noreply.github.com>
Date: Fri, 20 Dec 2024 08:35:57 +0100
Subject: [PATCH] feat(observe): add raw prompt for streamlit agent (#161)

Signed-off-by: Milan Gallas <Milan.Gallas@ibm.com>
Co-authored-by: Milan Gallas <Milan.Gallas@ibm.com>
---
 src/modules/chat/trace/useBuildTraceData.ts | 54 +++++++++------------
 src/modules/chat/trace/utils.ts             | 10 ++++
 2 files changed, 33 insertions(+), 31 deletions(-)

diff --git a/src/modules/chat/trace/useBuildTraceData.ts b/src/modules/chat/trace/useBuildTraceData.ts
index a9346e2..df51c4a 100644
--- a/src/modules/chat/trace/useBuildTraceData.ts
+++ b/src/modules/chat/trace/useBuildTraceData.ts
@@ -26,6 +26,7 @@ import {
   getLastNewTokenSpan,
   getGeneratedTokenCountSafe,
   getExecutionTime,
+  getRawPrompt,
 } from './utils';
 import { useEffect, useMemo, useState } from 'react';
 import { GENERATE_EVENT_TOOL_START, TraceData } from './types';
@@ -79,35 +80,24 @@ export function useBuildTraceData({ enabled, threadId, runId }: Props): {
     if (!spans) return null;
     const iterations = spans
       .filter((span) => span.name.includes('iteration-'))
-      .map((span) => {
-        return {
-          span,
-          children: loadNestedSpans(span, spans).flat(),
-        };
-      })
+      .map((span) => ({
+        span,
+        children: loadNestedSpans(span, spans).flat(),
+      }))
       .filter((span) => span.children.length >= 1)
-      .map((iteration) => {
-        const customEventData = iteration.children.find(
-          (span) => span.attributes.name === 'startCustom',
-        )?.attributes.data;
-        const rawPrompt = customEventData?.rawPrompt
-          ? String(customEventData.rawPrompt)
-          : undefined;
-
-        return {
-          groupId: iteration.span.name,
-          executionTime: getExecutionTime(iteration.children),
-          tokenCount: getGeneratedTokenCountSafe(
-            getLastNewTokenSpan(iteration.children),
-          ),
-          rawPrompt,
-          type: iteration.children.find(
-            (span) => span.attributes.name === GENERATE_EVENT_TOOL_START,
-          )
-            ? InterationType.TOOL
-            : InterationType.FINAL_ANSWER,
-        };
-      });
+      .map((iteration) => ({
+        groupId: iteration.span.name,
+        executionTime: getExecutionTime(iteration.children),
+        tokenCount: getGeneratedTokenCountSafe(
+          getLastNewTokenSpan(iteration.children),
+        ),
+        rawPrompt: getRawPrompt(iteration.children),
+        type: iteration.children.find(
+          (span) => span.attributes.name === GENERATE_EVENT_TOOL_START,
+        )
+          ? InterationType.TOOL
+          : InterationType.FINAL_ANSWER,
+      }));
 
     const executionTime = getExecutionTime(spans);
 
@@ -118,9 +108,11 @@ export function useBuildTraceData({ enabled, threadId, runId }: Props): {
         0,
       ) || getGeneratedTokenCountSafe(getLastNewTokenSpan(spans));
 
-    const rawPrompt = iterations.find(
-      (iteration) => iteration.type === InterationType.FINAL_ANSWER,
-    )?.rawPrompt;
+    // load rawPrompt from iterations if they exist. Otherwise, compute it from all spans.
+    const rawPrompt =
+      iterations.find(
+        (iteration) => iteration.type === InterationType.FINAL_ANSWER,
+      )?.rawPrompt || getRawPrompt(spans);
 
     return { executionTime, tokenCount, rawPrompt, iterations };
   }, [spans]);
diff --git a/src/modules/chat/trace/utils.ts b/src/modules/chat/trace/utils.ts
index 70aa337..d7c2b1a 100644
--- a/src/modules/chat/trace/utils.ts
+++ b/src/modules/chat/trace/utils.ts
@@ -51,3 +51,13 @@ export function getExecutionTime(spans: TraceSpan[]) {
     new Date(firstItem.start_time).getTime()
   );
 }
+
+export function getRawPrompt(spans: TraceSpan[]): string | undefined {
+  const customSpanData = spans.find(
+    (span) => span.attributes.name === 'startCustom',
+  )?.attributes.data;
+
+  return customSpanData?.rawPrompt
+    ? String(customSpanData.rawPrompt)
+    : undefined;
+}