Tool span + traces UI (lmnr-ai#371)

app-server/src/ch/spans.rs

@@ -29,6 +29,7 @@ impl Into<u8> for SpanType {
             SpanType::EXECUTOR => 3,
             SpanType::EVALUATOR => 4,
             SpanType::EVALUATION => 5,
+            SpanType::TOOL => 6,
         }
     }
 }

app-server/src/db/spans.rs

@@ -21,6 +21,7 @@ pub enum SpanType {
     EXECUTOR,
     EVALUATOR,
     EVALUATION,
+    TOOL,
 }
 
 impl FromStr for SpanType {
@@ -34,6 +35,7 @@ impl FromStr for SpanType {
             "EXECUTOR" => Ok(SpanType::EXECUTOR),
             "EVALUATOR" => Ok(SpanType::EVALUATOR),
             "EVALUATION" => Ok(SpanType::EVALUATION),
+            "TOOL" => Ok(SpanType::TOOL),
             _ => Err(anyhow::anyhow!("Invalid span type: {}", s)),
         }
     }
@@ -89,6 +91,7 @@ pub async fn record_span(pool: &PgPool, span: &Span, project_id: &Uuid) -> Resul
         ),
         &None => None,
     };
+
     sqlx::query(
         "INSERT INTO spans
             (span_id,

app-server/src/main.rs

@@ -79,7 +79,7 @@ mod traces;
 
 const DEFAULT_CACHE_SIZE: u64 = 100; // entries
 const HTTP_PAYLOAD_LIMIT: usize = 5 * 1024 * 1024; // 5MB
-const GRPC_PAYLOAD_DECODING_LIMIT: usize = 10 * 1024 * 1024; // 10MB
+const GRPC_PAYLOAD_DECODING_LIMIT: usize = 50 * 1024 * 1024; // 50MB
 
 fn tonic_error_to_io_error(err: tonic::transport::Error) -> io::Error {
     io::Error::new(io::ErrorKind::Other, err)

app-server/src/traces/span_attributes.rs

@@ -28,4 +28,5 @@ pub const GEN_AI_OUTPUT_COST: &str = "gen_ai.usage.output_cost";
 pub const ASSOCIATION_PROPERTIES_PREFIX: &str = "lmnr.association.properties";
 pub const SPAN_TYPE: &str = "lmnr.span.type";
 pub const SPAN_PATH: &str = "lmnr.span.path";
+pub const SPAN_IDS_PATH: &str = "lmnr.span.ids_path";
 pub const LLM_NODE_RENDERED_PROMPT: &str = "lmnr.span.prompt";

app-server/src/traces/spans.rs

@@ -27,9 +27,9 @@ use super::{
         ASSOCIATION_PROPERTIES_PREFIX, GEN_AI_COMPLETION_TOKENS, GEN_AI_INPUT_COST,
         GEN_AI_INPUT_TOKENS, GEN_AI_OUTPUT_COST, GEN_AI_OUTPUT_TOKENS, GEN_AI_PROMPT_TOKENS,
         GEN_AI_REQUEST_MODEL, GEN_AI_RESPONSE_MODEL, GEN_AI_SYSTEM, GEN_AI_TOTAL_COST,
-        GEN_AI_TOTAL_TOKENS, LLM_NODE_RENDERED_PROMPT, SPAN_PATH, SPAN_TYPE,
+        GEN_AI_TOTAL_TOKENS, LLM_NODE_RENDERED_PROMPT, SPAN_IDS_PATH, SPAN_PATH, SPAN_TYPE,
     },
-    utils::json_value_to_string,
+    utils::{json_value_to_string, skip_span_name},
 };
 
 const INPUT_ATTRIBUTE_NAME: &str = "lmnr.span.input";
@@ -174,6 +174,15 @@ impl SpanAttributes {
     }
 
     pub fn path(&self) -> Option<Vec<String>> {
+        let raw_path = self.raw_path();
+        raw_path.map(|path| {
+            path.into_iter()
+                .filter(|name| !skip_span_name(name))
+                .collect()
+        })
+    }
+
+    fn raw_path(&self) -> Option<Vec<String>> {
         match self.attributes.get(SPAN_PATH) {
             Some(Value::Array(arr)) => Some(
                 arr.iter()
@@ -189,6 +198,36 @@ impl SpanAttributes {
         self.path().map(|path| path.join("."))
     }
 
+    pub fn ids_path(&self) -> Option<Vec<String>> {
+        let attributes_ids_path = match self.attributes.get(SPAN_IDS_PATH) {
+            Some(Value::Array(arr)) => Some(
+                arr.iter()
+                    .map(|v| json_value_to_string(v.clone()))
+                    .collect::<Vec<_>>(),
+            ),
+            _ => None,
+        };
+
+        attributes_ids_path.map(|ids_path| {
+            let path = self.raw_path();
+            if let Some(path) = path {
+                ids_path
+                    .into_iter()
+                    .zip(path.into_iter())
+                    .filter_map(|(id, name)| {
+                        if skip_span_name(&name) {
+                            None
+                        } else {
+                            Some(id)
+                        }
+                    })
+                    .collect()
+            } else {
+                ids_path
+            }
+        })
+    }
+
     pub fn set_usage(&mut self, usage: &SpanUsage) {
         self.attributes
             .insert(GEN_AI_INPUT_TOKENS.to_string(), json!(usage.input_tokens));
@@ -240,6 +279,29 @@ impl SpanAttributes {
         }
     }
 
+    pub fn update_path(&mut self) {
+        self.attributes.insert(
+            SPAN_PATH.to_string(),
+            Value::Array(
+                self.path()
+                    .unwrap_or_default()
+                    .into_iter()
+                    .map(serde_json::Value::String)
+                    .collect(),
+            ),
+        );
+        self.attributes.insert(
+            SPAN_IDS_PATH.to_string(),
+            Value::Array(
+                self.ids_path()
+                    .unwrap_or_default()
+                    .into_iter()
+                    .map(serde_json::Value::String)
+                    .collect(),
+            ),
+        );
+    }
+
     pub fn labels(&self) -> HashMap<String, Value> {
         self.get_flattened_association_properties("label")
     }
@@ -309,6 +371,7 @@ impl Span {
 
     pub fn should_save(&self) -> bool {
         self.get_attributes().tracing_level() != Some(TracingLevel::Off)
+            && !skip_span_name(&self.name)
     }
 
     /// Create a span from an OpenTelemetry span.

app-server/src/traces/utils.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashMap, sync::Arc};
 
+use regex::Regex;
 use serde_json::Value;
 use uuid::Uuid;
 
@@ -116,6 +117,18 @@ pub async fn record_span_to_db(
     }
 
     span_attributes.extend_span_path(&span.name);
+    span_attributes.ids_path().map(|path| {
+        // set the parent to the second last id in the path
+        if path.len() > 1 {
+            let parent_id = path
+                .get(path.len() - 2)
+                .and_then(|id| Uuid::parse_str(id).ok());
+            if let Some(parent_id) = parent_id {
+                span.parent_span_id = Some(parent_id);
+            }
+        }
+    });
+    span_attributes.update_path();
     span.set_attributes(&span_attributes);
 
     let update_attrs_res =
@@ -178,3 +191,8 @@ pub async fn record_labels_to_db_and_ch(
 
     Ok(())
 }
+
+pub fn skip_span_name(name: &str) -> bool {
+    let re = Regex::new(r"^Runnable[A-Z][A-Za-z]*(?:<[A-Za-z_,]+>)*\.task$").unwrap();
+    re.is_match(name)
+}

frontend/components/traces/session-player.tsx

@@ -44,6 +44,7 @@ const SessionPlayer = forwardRef<SessionPlayerHandle, SessionPlayerProps>(
     const [dimensions, setDimensions] = useState({ width: 0, height: 0 });
     const [startTime, setStartTime] = useState(0);
     const { projectId } = useProjectContext();
+    const workerRef = useRef<Worker | null>(null);
 
     // Add resize observer effect
     useEffect(() => {
@@ -64,23 +65,37 @@ const SessionPlayer = forwardRef<SessionPlayerHandle, SessionPlayerProps>(
     // Add debounce timer ref
     const debounceTimerRef = useRef<NodeJS.Timeout | null>(null);
 
-    // Create debounced goto function
-    const debouncedGoto = useCallback((time: number) => {
-      if (debounceTimerRef.current) {
-        clearTimeout(debounceTimerRef.current);
-      }
+    // Initialize worker
+    useEffect(() => {
+      workerRef.current = new Worker(new URL('@/lib/workers/player-worker.ts', import.meta.url));
 
-      debounceTimerRef.current = setTimeout(() => {
+      workerRef.current.onmessage = (e) => {
+        const { result, isPlaying } = e.data;
         if (playerRef.current) {
           try {
-            playerRef.current.goto(time * 1000, isPlaying);
+            playerRef.current.goto(result, isPlaying);
           } catch (e) {
             console.error(e);
           }
         }
-      }, 50); // 50ms debounce delay
+      };
+
+      return () => {
+        workerRef.current?.terminate();
+      };
     }, []);
 
+    // Update the debouncedGoto function
+    const debouncedGoto = useCallback((time: number) => {
+      if (debounceTimerRef.current) {
+        clearTimeout(debounceTimerRef.current);
+      }
+
+      debounceTimerRef.current = setTimeout(() => {
+        workerRef.current?.postMessage({ time, isPlaying });
+      }, 50);
+    }, [isPlaying]);
+
     const getEvents = async () => {
       const res = await fetch(`/api/projects/${projectId}/browser-sessions/events?traceId=${traceId}`, {
         method: 'GET',
@@ -109,12 +124,17 @@ const SessionPlayer = forwardRef<SessionPlayerHandle, SessionPlayerProps>(
 
     useEffect(() => {
       if (hasBrowserSession) {
+        setEvents([]);
+        setIsPlaying(false);
+        setCurrentTime(0);
+        setTotalDuration(0);
+        setSpeed(1);
         getEvents();
       }
-    }, [hasBrowserSession]);
+    }, [hasBrowserSession, traceId]);
 
     useEffect(() => {
-      if (!events?.length || !playerContainerRef.current || playerRef.current) return;
+      if (!events?.length || !playerContainerRef.current) return;
 
       try {
         playerRef.current = new rrwebPlayer({
@@ -197,16 +217,15 @@ const SessionPlayer = forwardRef<SessionPlayerHandle, SessionPlayerProps>(
       setSpeed((currentSpeed) => (currentSpeed === 1 ? 2 : 1));
     };
 
-    // Expose imperative methods to parent
+    // Update the imperative handle
     useImperativeHandle(ref, () => ({
       goto: (time: number) => {
         if (playerRef.current) {
-          playerRef.current.pause();
-          playerRef.current.goto(time * 1000, isPlaying);
           setCurrentTime(time);
+          workerRef.current?.postMessage({ time, isPlaying });
         }
       }
-    }), []);
+    }), [isPlaying]);
 
     return (
       <>

frontend/components/traces/span-card.tsx

@@ -5,7 +5,6 @@ import { getDuration, getDurationString } from '@/lib/flow/utils';
 import { Span } from '@/lib/traces/types';
 import { cn, formatSecondsToMinutesAndSeconds } from '@/lib/utils';
 
-import { Label } from '../ui/label';
 import SpanTypeIcon from './span-type-icon';
 
 const ROW_HEIGHT = 36;
@@ -90,9 +89,9 @@ export function SpanCard({
           <div className="text-ellipsis overflow-hidden whitespace-nowrap text-base truncate max-w-[150px]">
             {span.name}
           </div>
-          <Label className="text-secondary-foreground">
+          <div className="text-secondary-foreground px-2 py-0.5 bg-secondary rounded-full text-xs">
             {getDurationString(span.startTime, span.endTime)}
-          </Label>
+          </div>
           <div
             className="z-30 top-[-px]  hover:bg-red-100/10 absolute transition-all"
             style={{

frontend/components/traces/span-type-icon.tsx

@@ -1,4 +1,4 @@
-import { Activity, ArrowRight, Braces, Gauge, MessageCircleMore } from "lucide-react";
+import { Activity, ArrowRight, Bolt, Braces, Gauge, MessageCircleMore } from "lucide-react";
 
 import { SpanType } from "@/lib/traces/types";
 import { SPAN_TYPE_TO_COLOR } from "@/lib/traces/utils";
@@ -46,6 +46,9 @@ export default function SpanTypeIcon({
       {spanType === SpanType.EVALUATION && (
         <Gauge size={size} />
       )}
+      {spanType === SpanType.TOOL && (
+        <Bolt size={size} />
+      )}
     </div>
   );
 }

frontend/components/traces/span-view.tsx

@@ -1,15 +1,10 @@
-import {
-  Activity,
-  ArrowRight,
-  Braces,
-  Gauge,
-  MessageCircleMore,
-} from 'lucide-react';
+
+
 import useSWR from 'swr';
 
 import { useProjectContext } from '@/contexts/project-context';
 import { Event } from '@/lib/events/types';
-import { Span, SpanType } from '@/lib/traces/types';
+import { Span } from '@/lib/traces/types';
 import { swrFetcher } from '@/lib/utils';
 
 import Formatter from '../ui/formatter';
@@ -19,6 +14,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from '../ui/tabs';
 import { AddLabelPopover } from './add-label-popover';
 import AddToLabelingQueuePopover from './add-to-labeling-queue-popover';
 import ExportSpansDialog from './export-spans-dialog';
+import SpanTypeIcon from './span-type-icon';
 import { SpanViewSpan } from './span-view-span';
 import StatsShields from './stats-shields';
 
@@ -55,20 +51,10 @@ export function SpanView({ spanId }: SpanViewProps) {
     <>
       <Tabs className="flex flex-col h-full w-full" defaultValue="span">
         <div className="border-b flex-none">
-          <div className="flex flex-col px-4 pt-2 gap-2">
-            <div className='flex flex-col gap-2'>
+          <div className="flex flex-col px-4 pt-2 gap-1">
+            <div className='flex flex-col gap-1'>
               <div className="flex flex-none items-center space-x-2">
-                <div className="p-1.5 px-2 text-xs text-secondary-foreground rounded bg-secondary">
-                  {span.spanType === SpanType.DEFAULT && <Braces size={16} />}
-                  {span.spanType === SpanType.LLM && (
-                    <MessageCircleMore size={16} />
-                  )}
-                  {span.spanType === SpanType.EXECUTOR && <Activity size={16} />}
-                  {span.spanType === SpanType.EVALUATOR && (
-                    <ArrowRight size={16} />
-                  )}
-                  {span.spanType === SpanType.EVALUATION && <Gauge size={16} />}
-                </div>
+                <SpanTypeIcon spanType={span.spanType} />
                 <div className="flex-grow text-xl items-center font-medium truncate max-w-[400px]">
                   {span.name}
                 </div>

frontend/components/traces/timeline.tsx

@@ -80,6 +80,10 @@ export default function Timeline({
       traverse(span, childSpans, orderedSpans);
     }
 
+    if (orderedSpans.length === 0) {
+      return;
+    }
+
     let startTime = new Date(orderedSpans[0].startTime).getTime();
     let endTime = new Date(orderedSpans[orderedSpans.length - 1].endTime).getTime();