Merge pull request #625 from lmnr-ai/dev

langgraph, playground enhancements, tracing enhancements, evals UI enhancements

Author: dinmukhamedm

app-server/Cargo.lock

@@ -56,6 +56,7 @@ tonic = {version = "0.13", features = ["gzip"]}
 url = "2.5.4"
 uuid = {version = "1.16.0", features = ["v4", "fast-rng", "macro-diagnostics", "serde"]}
 reqwest = { version = "0.12.15", features = ["json"] }
+itertools = "0.14.0"
 
 [build-dependencies]
 tonic-build = "0.13"

app-server/Cargo.toml

@@ -24,7 +24,7 @@ pub struct EvaluationDatapointResult {
     #[serde(default)]
     pub target: Value,
     #[serde(default)]
-    pub metadata: HashMap<String, Value>,
+    pub metadata: Option<HashMap<String, Value>>,
     #[serde(default)]
     pub executor_output: Option<Value>,
     #[serde(default)]
@@ -62,7 +62,7 @@ pub fn get_columns_from_points(points: &Vec<EvaluationDatapointResult>) -> Datap
 
     let metadatas = points
         .iter()
-        .map(|point| point.metadata.clone())
+        .map(|point| point.metadata.clone().unwrap_or_default())
         .collect::<Vec<_>>();
 
     let executor_outputs = points

app-server/src/evaluations/utils.rs

@@ -1,10 +1,10 @@
 use std::{collections::HashMap, sync::Arc};
 
 use crate::{
-    cache::{keys::LLM_PRICES_CACHE_KEY, Cache, CacheTrait},
+    cache::{Cache, CacheTrait, keys::LLM_PRICES_CACHE_KEY},
     db::{
-        prices::{get_price, DBPriceEntry},
         DB,
+        prices::{DBPriceEntry, get_price},
     },
     traces::spans::InputTokens,
 };
@@ -47,7 +47,7 @@ pub async fn estimate_output_cost(
     let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await.ok()?;
 
     let price_per_million_tokens = match cache_res {
-        Some(price) => price.input_price_per_million,
+        Some(price) => price.output_price_per_million,
         None => {
             let price = get_price(&db.pool, provider, model).await.ok()?;
             let price = LLMPriceEntry::from(price);

app-server/src/language_model/costs/mod.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use itertools::Itertools;
 use utils::prepare_span_for_recording;
 use uuid::Uuid;
 
@@ -50,6 +51,28 @@ pub async fn process_spans_and_events(
     let span_usage =
         get_llm_usage_for_span(&mut span.get_attributes(), db.clone(), cache.clone()).await;
 
+    let mut has_seen_first_token = false;
+
+    // OpenLLMetry auto-instrumentation sends this event for every chunk
+    // While this is helpful to get TTFT, we don't want to store excessive,
+    // so we only keep the first one.
+    let events = events
+        .into_iter()
+        .sorted_by(|a, b| a.timestamp.cmp(&b.timestamp))
+        .filter(|event| {
+            if event.name == "llm.content.completion.chunk" {
+                if !has_seen_first_token {
+                    has_seen_first_token = true;
+                    true
+                } else {
+                    false
+                }
+            } else {
+                true
+            }
+        })
+        .collect();
+
     let trace_attributes = prepare_span_for_recording(span, &span_usage, &events);
 
     if let Some(span_path) = span.get_attributes().path() {

app-server/src/traces/mod.rs

@@ -31,16 +31,7 @@ pub async fn push_spans_to_queue(
                 let events = otel_span
                     .events
                     .into_iter()
-                    .filter_map(|event| {
-                        // OpenLLMetry auto-instrumentation sends this event for every chunk
-                        // While this is helpful to get TTFT, we don't want to store excessive
-                        // events
-                        if event.name == "llm.content.completion.chunk" {
-                            None
-                        } else {
-                            Some(Event::from_otel(event, span.span_id, project_id))
-                        }
-                    })
+                    .map(|event| Event::from_otel(event, span.span_id, project_id))
                     .collect::<Vec<Event>>();
 
                 if !span.should_save() {

app-server/src/traces/producer.rs

@@ -935,8 +935,12 @@ fn output_message_from_completion_content(
         }
     } else {
         let mut out_vec = if let Some(Value::String(s)) = msg_content {
-            let text_block = ChatMessageContentPart::Text(ChatMessageText { text: s.clone() });
-            vec![serde_json::to_value(text_block).unwrap()]
+            if s.is_empty() {
+                vec![]
+            } else {
+                let text_block = ChatMessageContentPart::Text(ChatMessageText { text: s.clone() });
+                vec![serde_json::to_value(text_block).unwrap()]
+            }
         } else {
             vec![]
         };

app-server/src/traces/spans.rs

@@ -0,0 +1,18 @@
+import { openai } from "@ai-sdk/openai";
+import { generateText } from "ai";
+
+export async function POST(req: Request) {
+  const { prompt }: { prompt: string } = await req.json();
+
+  if (!process.env.OPENAI_API_KEY) {
+    return Response.json({ error: "OPENAI_API_KEY is not set" }, { status: 500 });
+  }
+
+  const { text } = await generateText({
+    model: openai("gpt-4.1-nano"),
+    system: "You are a helpful assistant.",
+    prompt,
+  });
+
+  return Response.json({ text });
+}

frontend/app/api/completion/route.ts

@@ -1,6 +1,7 @@
 import { and, asc, eq, inArray, sql } from "drizzle-orm";
 import { NextRequest } from "next/server";
 
+import { DatatableFilter } from "@/components/ui/datatable-filter/utils";
 import { searchSpans } from "@/lib/clickhouse/spans";
 import { SpanSearchType } from "@/lib/clickhouse/types";
 import { db } from "@/lib/db/drizzle";
@@ -10,7 +11,6 @@ import {
   EvaluationScoreDistributionBucket,
   EvaluationScoreStatistics,
 } from "@/lib/evaluation/types";
-import { DatatableFilter } from "@/lib/types";
 
 // Constants for distribution calculation
 const DEFAULT_LOWER_BOUND = 0.0;

frontend/app/api/projects/[projectId]/evaluations/[evaluationId]/route.ts

@@ -1,11 +1,14 @@
-import { and, asc, eq, inArray, sql } from "drizzle-orm";
+import { and, asc, eq, inArray, not, sql } from "drizzle-orm";
+import { partition } from "lodash";
 import { NextRequest, NextResponse } from "next/server";
 
 import { searchSpans } from "@/lib/clickhouse/spans";
 import { SpanSearchType } from "@/lib/clickhouse/types";
 import { TimeRange } from "@/lib/clickhouse/utils";
 import { db } from "@/lib/db/drizzle";
-import { events, spans } from "@/lib/db/migrations/schema";
+import { events, labelClasses, labels, spans } from "@/lib/db/migrations/schema";
+import { FilterDef, filtersToSql } from "@/lib/db/modifiers";
+
 export async function GET(
   req: NextRequest,
   props: { params: Promise<{ projectId: string; traceId: string }> }
@@ -15,6 +18,63 @@ export async function GET(
   const traceId = params.traceId;
   const searchQuery = req.nextUrl.searchParams.get("search");
   const searchType = req.nextUrl.searchParams.getAll("searchIn");
+
+  const urlParamFilters = (() => {
+    try {
+      const rawFilters = req.nextUrl.searchParams.getAll("filter").map((f) => JSON.parse(f) as FilterDef);
+      return Array.isArray(rawFilters) ? rawFilters : [];
+    } catch {
+      return [];
+    }
+  })();
+
+  const [filters, statusFilters] = partition(urlParamFilters, (f) => f.column !== "status");
+  const [otherFilters, tagsFilters] = partition(filters, (f) => f.column !== "tags");
+
+  const statusSqlFilters = statusFilters.map((filter) => {
+    if (filter.value === "success") {
+      return filter.operator === "eq" ? sql`status IS NULL` : sql`status IS NOT NULL`;
+    } else if (filter.value === "error") {
+      return filter.operator === "eq" ? sql`status = 'error'` : sql`status != 'error' OR status IS NULL`;
+    }
+    return sql`1=1`;
+  });
+
+  const tagsSqlFilters = tagsFilters.map((filter) => {
+    const name = filter.value;
+    const inArrayFilter = inArray(
+      spans.spanId,
+      db
+        .select({ span_id: spans.spanId })
+        .from(spans)
+        .innerJoin(labels, eq(spans.spanId, labels.spanId))
+        .innerJoin(labelClasses, eq(labels.classId, labelClasses.id))
+        .where(and(eq(labelClasses.name, name)))
+    );
+    return filter.operator === "eq" ? inArrayFilter : not(inArrayFilter);
+  });
+
+  const processedFilters = otherFilters.map((filter) => {
+    if (filter.column === "path") {
+      filter.column = "(attributes ->> 'lmnr.span.path')";
+    } else if (filter.column === "tokens") {
+      filter.column = "(attributes ->> 'llm.usage.total_tokens')::int8";
+    } else if (filter.column === "cost") {
+      filter.column = "(attributes ->> 'gen_ai.usage.cost')::float8";
+    } else if (filter.column === "model") {
+      filter.column = "(attributes ->> 'gen_ai.request.model')";
+    }
+    return filter;
+  });
+
+  const sqlFilters = filtersToSql(
+    processedFilters,
+    [new RegExp(/^\(attributes\s*->>\s*'[a-zA-Z_\.]+'\)(?:::int8|::float8)?$/)],
+    {
+      latency: sql<number>`EXTRACT(EPOCH FROM (end_time - start_time))`,
+    }
+  );
+
   let searchSpanIds: string[] = [];
   if (searchQuery) {
     const timeRange = { pastHours: "all" } as TimeRange;
@@ -32,7 +92,7 @@ export async function GET(
   const spanEventsQuery = db.$with("span_events").as(
     db
       .select({
-        spanId: events.spanId,
+        eventSpanId: sql`events.span_id`.as("eventSpanId"),
         projectId: events.projectId,
         events: sql`jsonb_agg(jsonb_build_object(
         'id', events.id,
@@ -78,12 +138,15 @@ export async function GET(
     .from(spans)
     .leftJoin(
       spanEventsQuery,
-      and(eq(spans.spanId, spanEventsQuery.spanId), eq(spans.projectId, spanEventsQuery.projectId))
+      and(eq(spans.spanId, spanEventsQuery.eventSpanId), eq(spans.projectId, spanEventsQuery.projectId))
     )
     .where(
       and(
         eq(spans.traceId, traceId),
         eq(spans.projectId, projectId),
+        ...sqlFilters,
+        ...statusSqlFilters,
+        ...tagsSqlFilters,
         ...(searchQuery !== null ? [inArray(spans.spanId, searchSpanIds)] : [])
       )
     )
@@ -94,7 +157,7 @@ export async function GET(
   return NextResponse.json(
     spanItems.map((span) => ({
       ...span,
-      parentSpanId: searchSpanIds.length > 0 ? null : span.parentSpanId,
+      parentSpanId: searchSpanIds.length > 0 || urlParamFilters.length > 0 ? null : span.parentSpanId,
     }))
   );
 }