Updates all GenAI insights to filter to those created by automatic LLM

instrumentation
DAGWorks-Inc · Aug 27, 2024 · 9244706 · 9244706
1 parent b9af8b0
commit 9244706
Showing 1 changed file with 22 additions and 8 deletions.
diff --git a/telemetry/ui/src/components/routes/app/InsightsView.tsx b/telemetry/ui/src/components/routes/app/InsightsView.tsx
@@ -54,20 +54,25 @@ const REGISTERED_INSIGHTS: Insight[] = [
   {
     category: 'llm',
     hasInsight: (allAttributes) => {
-      return allAttributes.some((attribute) => attribute.key.endsWith('prompt_tokens'));
+      return allAttributes.some(
+        (attribute) => attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')
+      );
     },
     insightName: 'Total Prompt Tokens',
     RenderInsightValue: (props) => {
       let totalPromptTokens = 0;
       props.attributes.forEach((attribute) => {
-        if (attribute.key.endsWith('prompt_tokens')) {
+        if (attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')) {
           totalPromptTokens += attribute.value as number;
         }
       });
       return <p>{totalPromptTokens}</p>;
     },
     captureIndividualValues: (allAttributes) => {
-      return allAttributes.filter((attribute) => attribute.key.endsWith('prompt_tokens'));
+      console.log(allAttributes.map((i) => i.key));
+      return allAttributes.filter(
+        (attribute) => attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')
+      );
     },
     RenderIndividualValue: (props: { attribute: AttributeModel }) => {
       return <p>{props.attribute.value?.toString()}</p>;
@@ -76,7 +81,10 @@ const REGISTERED_INSIGHTS: Insight[] = [
   {
     category: 'llm',
     hasInsight: (allAttributes) => {
-      return allAttributes.some((attribute) => attribute.key.endsWith('completion_tokens'));
+      return allAttributes.some(
+        (attribute) =>
+          attribute.key.endsWith('completion_tokens') && attribute.key.startsWith('gen_ai')
+      );
     },
     insightName: 'Total Completion Tokens',
     RenderInsightValue: (props) => {
@@ -89,7 +97,10 @@ const REGISTERED_INSIGHTS: Insight[] = [
       return <p>{totalCompletionTokens}</p>;
     },
     captureIndividualValues: (allAttributes) => {
-      return allAttributes.filter((attribute) => attribute.key.endsWith('completion_tokens'));
+      return allAttributes.filter(
+        (attribute) =>
+          attribute.key.endsWith('completion_tokens') && attribute.key.startsWith('gen_ai')
+      );
     },
     RenderIndividualValue: (props: { attribute: AttributeModel }) => {
       return <p>{props.attribute.value?.toString()}</p>;
@@ -98,13 +109,15 @@ const REGISTERED_INSIGHTS: Insight[] = [
   {
     category: 'llm',
     hasInsight: (allAttributes) => {
-      return allAttributes.some((attribute) => attribute.key.endsWith('prompt_tokens'));
+      return allAttributes.some(
+        (attribute) => attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')
+      );
     },
     insightName: 'Total LLM Calls',
     RenderInsightValue: (props) => {
       let totalLLMCalls = 0;
       props.attributes.forEach((attribute) => {
-        if (attribute.key.endsWith('prompt_tokens')) {
+        if (attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')) {
           totalLLMCalls += 1;
         }
       });
@@ -113,7 +126,7 @@ const REGISTERED_INSIGHTS: Insight[] = [
     captureIndividualValues: (allAttributes) => {
       const spanIDToLLMCalls = new Map<string, number>();
       allAttributes.forEach((attribute) => {
-        if (attribute.key.endsWith('prompt_tokens')) {
+        if (attribute.key.endsWith('prompt_tokens') && attribute.key.startsWith('gen_ai')) {
           spanIDToLLMCalls.set(
             attribute.span_id || '',
             (spanIDToLLMCalls.get(attribute.span_id || '') || 0) + 1
@@ -160,6 +173,7 @@ const InsightSubTable = (props: {
 
   const { currentSelectedIndex, setCurrentSelectedIndex, currentHoverIndex, setCurrentHoverIndex } =
     useContext(AppViewHighlightContext);
+  console.log(individualValues);
 
   return (
     <>