improve llm spans view

frontend/components/shared/traces/span-view.tsx

@@ -62,7 +62,7 @@ export function SpanView({ spanId, traceId }: SpanViewProps) {
               <MonoWithCopy className="text-muted-foreground">{span.spanId}</MonoWithCopy>
             </div>
             <div className="flex flex-wrap py-1 gap-2">
-              <SpanStatsShields startTime={span.startTime} endTime={span.endTime} attributes={span.attributes}>
+              <SpanStatsShields span={span}>
                 <div className="flex flex-row text-xs font-mono space-x-2 rounded-md p-0.5 px-2 border items-center">
                   {new Date(span.startTime).toLocaleString()}
                 </div>

frontend/components/shared/traces/trace-view.tsx

@@ -86,7 +86,6 @@ const PureTraceView = ({ trace, spans }: TraceViewProps) => {
     spanPath: state.spanPath,
     setSpanPath: state.setSpanPath,
   }));
-
   const hasLangGraph = useMemo(() => getHasLangGraph(), [getHasLangGraph]);
   const llmSpanIds = useMemo(
     () =>

frontend/components/traces/span-controls.tsx

@@ -1,5 +1,5 @@
 import { get } from "lodash";
-import {ChevronDown, Copy, Database, Loader, PlayCircle} from "lucide-react";
+import { ChevronDown, Copy, Database, Loader, PlayCircle } from "lucide-react";
 import Link from "next/link";
 import { useParams } from "next/navigation";
 import React, { PropsWithChildren, useCallback, useMemo } from "react";
@@ -14,7 +14,12 @@ import ErrorCard from "@/components/traces/error-card";
 import ExportSpansPopover from "@/components/traces/export-spans-popover";
 import { useOpenInSql } from "@/components/traces/trace-view/use-open-in-sql.tsx";
 import { Button } from "@/components/ui/button";
-import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from "@/components/ui/dropdown-menu";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
 import { Event } from "@/lib/events/types";
 import { useToast } from "@/lib/hooks/use-toast";
 import { Span, SpanType } from "@/lib/traces/types";
@@ -37,7 +42,10 @@ export function SpanControls({ children, span, events }: PropsWithChildren<SpanC
   );
 
   const { toast } = useToast();
-  const { openInSql, isLoading } = useOpenInSql({ projectId: projectId as string, params: { type: 'span', spanId: span.spanId } });
+  const { openInSql, isLoading } = useOpenInSql({
+    projectId: projectId as string,
+    params: { type: "span", spanId: span.spanId },
+  });
 
   const handleCopySpanId = useCallback(async () => {
     if (span?.spanId) {
@@ -82,12 +90,7 @@ export function SpanControls({ children, span, events }: PropsWithChildren<SpanC
           )}
         </div>
         <div className="flex flex-col flex-wrap gap-1.5">
-          <SpanStatsShields
-            className="flex-wrap"
-            startTime={span.startTime}
-            endTime={span.endTime}
-            attributes={span.attributes}
-          >
+          <SpanStatsShields className="flex-wrap" span={span}>
             <div className="text-xs font-mono space-x-2 rounded-md p-0.5 truncate px-2 border items-center">
               {new Date(span.startTime).toLocaleString()}
             </div>
@@ -101,7 +104,6 @@ export function SpanControls({ children, span, events }: PropsWithChildren<SpanC
             </div>
             <TagsList />
             <EvaluatorScoresList spanId={span.spanId} />
-
           </TagsContextProvider>
         </div>
 

frontend/components/traces/stats-shields.tsx

@@ -1,12 +1,13 @@
 import { TooltipPortal } from "@radix-ui/react-tooltip";
-import { compact, get, isNil, sortBy, uniq } from "lodash";
+import { compact, get, isNil, pick, sortBy, uniq } from "lodash";
 import { Bolt, Braces, ChevronDown, CircleDollarSign, Clock3, Coins } from "lucide-react";
 import { memo, PropsWithChildren } from "react";
 
-import { TraceViewTrace } from "@/components/traces/trace-view/trace-view-store.tsx";
+import { TraceViewSpan, TraceViewTrace } from "@/components/traces/trace-view/trace-view-store.tsx";
 import { DropdownMenu, DropdownMenuContent, DropdownMenuTrigger } from "@/components/ui/dropdown-menu";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
+import { Span } from "@/lib/traces/types.ts";
 import { cn, getDurationString, pluralize } from "@/lib/utils";
 
 import ContentRenderer from "../ui/content-renderer/index";
@@ -18,9 +19,7 @@ interface TraceStatsShieldsProps {
 }
 
 interface SpanStatsShieldsProps {
-  startTime: string;
-  endTime: string;
-  attributes: Record<string, any>;
+  span: Span;
   className?: string;
 }
 
@@ -119,61 +118,48 @@ const extractToolsFromAttributes = (attributes: Record<string, any>): Tool[] =>
       const name = attributes[`llm.request.functions.${index}.name`];
       const description = attributes[`llm.request.functions.${index}.description`];
       const rawParameters = attributes[`llm.request.functions.${index}.parameters`];
-      const parameters = typeof rawParameters === "string"
-        ? rawParameters
-        : JSON.stringify(rawParameters || {});
+      const parameters = typeof rawParameters === "string" ? rawParameters : JSON.stringify(rawParameters || {});
 
       return name ? { name, description, parameters } : null;
     })
   );
 };
 
 function StatsShieldsContent({
-  startTime,
-  endTime,
-  totalTokens,
-  inputTokens,
-  outputTokens,
-  inputCost,
-  outputCost,
-  totalCost,
+  stats,
   className,
   children,
 }: PropsWithChildren<{
-  startTime: string;
-  endTime: string;
-  totalTokens: number;
-  inputTokens: number;
-  outputTokens: number;
-  inputCost: number;
-  outputCost: number;
-  totalCost: number;
+  stats: Pick<
+    TraceViewSpan,
+    "startTime" | "endTime" | "inputTokens" | "outputTokens" | "totalTokens" | "inputCost" | "outputCost" | "totalCost"
+  >;
   className?: string;
 }>) {
   return (
     <div className={cn("flex items-center gap-2 font-mono min-w-0", className)}>
       <div className="flex space-x-1 items-center p-0.5 min-w-8 px-2 border rounded-md">
         <Clock3 size={12} className="min-w-3 min-h-3" />
-        <Label className="text-xs truncate text-foreground" title={getDurationString(startTime, endTime)}>
-          {getDurationString(startTime, endTime)}
+        <Label className="text-xs truncate text-foreground" title={getDurationString(stats.startTime, stats.endTime)}>
+          {getDurationString(stats.startTime, stats.endTime)}
         </Label>
       </div>
       <TooltipProvider delayDuration={250}>
         <Tooltip>
           <TooltipTrigger className="min-w-8">
             <div className="flex space-x-1 items-center p-0.5 min-w-8 px-2 border rounded-md">
               <Coins className="min-w-3" size={12} />
-              <Label className="text-xs truncate text-foreground">{totalTokens}</Label>
+              <Label className="text-xs truncate text-foreground">{stats.totalTokens}</Label>
             </div>
           </TooltipTrigger>
           <TooltipPortal>
             <TooltipContent side="bottom" className="p-2 border">
               <div className="flex-col space-y-1">
                 <Label className="flex text-xs gap-1">
-                  <span className="text-secondary-foreground">Input tokens</span> {inputTokens}
+                  <span className="text-secondary-foreground">Input tokens</span> {stats.inputTokens}
                 </Label>
                 <Label className="flex text-xs gap-1">
-                  <span className="text-secondary-foreground">Output tokens</span> {outputTokens}
+                  <span className="text-secondary-foreground">Output tokens</span> {stats.outputTokens}
                 </Label>
               </div>
             </TooltipContent>
@@ -185,20 +171,20 @@ function StatsShieldsContent({
           <TooltipTrigger className="min-w-8">
             <div className="flex space-x-1 items-center p-0.5 px-2 min-w-8 border rounded-md">
               <CircleDollarSign className="min-w-3" size={12} />
-              <Label className="text-xs truncate text-foreground">{totalCost?.toFixed(3)}</Label>
+              <Label className="text-xs truncate text-foreground">{stats.totalCost?.toFixed(3)}</Label>
             </div>
           </TooltipTrigger>
           <TooltipPortal>
             <TooltipContent side="bottom" className="p-2 border">
               <div className="flex-col space-y-1">
                 <Label className="flex text-xs gap-1">
-                  <span className="text-secondary-foreground">Total cost</span> {"$" + totalCost?.toFixed(5)}
+                  <span className="text-secondary-foreground">Total cost</span> {"$" + stats.totalCost?.toFixed(5)}
                 </Label>
                 <Label className="flex text-xs gap-1">
-                  <span className="text-secondary-foreground">Input cost</span> {"$" + inputCost?.toFixed(5)}
+                  <span className="text-secondary-foreground">Input cost</span> {"$" + stats.inputCost?.toFixed(5)}
                 </Label>
                 <Label className="flex text-xs gap-1">
-                  <span className="text-secondary-foreground">Output cost</span> {"$" + outputCost?.toFixed(5)}
+                  <span className="text-secondary-foreground">Output cost</span> {"$" + stats.outputCost?.toFixed(5)}
                 </Label>
               </div>
             </TooltipContent>
@@ -212,49 +198,41 @@ function StatsShieldsContent({
 
 const PureTraceStatsShields = ({ trace, className, children }: PropsWithChildren<TraceStatsShieldsProps>) => (
   <StatsShieldsContent
-    startTime={trace.startTime}
-    endTime={trace.endTime}
-    totalTokens={trace.totalTokens}
-    inputTokens={trace.inputTokens}
-    outputTokens={trace.outputTokens}
-    inputCost={trace.inputCost}
-    outputCost={trace.outputCost}
-    totalCost={trace.totalCost}
+    stats={pick(trace, [
+      "startTime",
+      "endTime",
+      "inputTokens",
+      "outputTokens",
+      "totalTokens",
+      "inputCost",
+      "outputCost",
+      "totalCost",
+    ])}
     className={className}
   >
     {children}
   </StatsShieldsContent>
 );
 
-const SpanStatsShields = ({
-  startTime,
-  endTime,
-  attributes,
-  className,
-  children,
-}: PropsWithChildren<SpanStatsShieldsProps>) => {
-  const inputTokenCount = get(attributes, "gen_ai.usage.input_tokens", 0);
-  const outputTokenCount = get(attributes, "gen_ai.usage.output_tokens", 0);
-  const totalTokenCount = inputTokenCount + outputTokenCount;
-  const inputCost = get(attributes, "gen_ai.usage.input_cost", 0);
-  const outputCost = get(attributes, "gen_ai.usage.output_cost", 0);
-  const cost = get(attributes, "gen_ai.usage.cost", 0);
-  const model = get(attributes, "gen_ai.response.model") || get(attributes, "gen_ai.request.model") || "";
-  const tools = extractToolsFromAttributes(attributes);
+const SpanStatsShields = ({ span, className, children }: PropsWithChildren<SpanStatsShieldsProps>) => {
+  const model = get(span.attributes, "gen_ai.response.model") || get(span.attributes, "gen_ai.request.model") || "";
+  const tools = extractToolsFromAttributes(span.attributes);
   const structuredOutputSchema =
-    get(attributes, "gen_ai.request.structured_output_schema") || get(attributes, "ai.schema");
+    get(span.attributes, "gen_ai.request.structured_output_schema") || get(span.attributes, "ai.schema");
 
   return (
     <div className="flex flex-wrap flex-col gap-1.5">
       <StatsShieldsContent
-        startTime={startTime}
-        endTime={endTime}
-        totalTokens={totalTokenCount}
-        inputTokens={inputTokenCount}
-        outputTokens={outputTokenCount}
-        inputCost={inputCost}
-        outputCost={outputCost}
-        totalCost={cost}
+        stats={pick(span, [
+          "startTime",
+          "endTime",
+          "inputTokens",
+          "outputTokens",
+          "totalTokens",
+          "inputCost",
+          "outputCost",
+          "totalCost",
+        ])}
         className={className}
       >
         {children}