From 5258fea7d5f4b3418fcc37b739df466f924ce8c9 Mon Sep 17 00:00:00 2001
From: Din <dinmukhamed.mailibay@gmail.com>
Date: Mon, 4 Nov 2024 19:05:18 -0800
Subject: [PATCH] overwrite span id when attribute set, query unique traces

---
 app-server/src/db/trace.rs     |  3 ++-
 app-server/src/traces/spans.rs | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/app-server/src/db/trace.rs b/app-server/src/db/trace.rs
index 3602d049..a9ec1ee4 100644
--- a/app-server/src/db/trace.rs
+++ b/app-server/src/db/trace.rs
@@ -402,6 +402,7 @@ pub async fn get_traces(
     query.push(
         "
         SELECT
+            DISTINCT ON (start_time, id)
             id,
             start_time,
             end_time,
@@ -435,7 +436,7 @@ pub async fn get_traces(
     add_filters_to_traces_query(&mut query, &filters);
 
     query
-        .push(" ORDER BY start_time DESC OFFSET ")
+        .push(" ORDER BY start_time DESC, id OFFSET ")
         .push_bind(offset as i64)
         .push(" LIMIT ")
         .push_bind(limit as i64);
diff --git a/app-server/src/traces/spans.rs b/app-server/src/traces/spans.rs
index 8818d9fc..b1b26e29 100644
--- a/app-server/src/traces/spans.rs
+++ b/app-server/src/traces/spans.rs
@@ -31,6 +31,10 @@ use super::span_attributes::{
 
 const INPUT_ATTRIBUTE_NAME: &str = "lmnr.span.input";
 const OUTPUT_ATTRIBUTE_NAME: &str = "lmnr.span.output";
+/// If this attribute is set to true, the parent span will be overridden with
+/// null. We hackily use this when we wrap a span in a NonRecordingSpan that
+/// is not sent to the backend – this is done to overwrite trace IDs for spans.
+const OVERRIDE_PARENT_SPAN_ATTRIBUTE_NAME: &str = "lmnr.internal.override_parent_span";
 
 pub struct SpanAttributes {
     pub attributes: HashMap<String, Value>,
@@ -309,6 +313,12 @@ impl Span {
             }
         }
 
+        // Spans with this attribute are wrapped in a NonRecordingSpan that, and we only
+        // do that when we add a new span to a trace as a root span.
+        if let Some(Value::Bool(true)) = attributes.get(OVERRIDE_PARENT_SPAN_ATTRIBUTE_NAME) {
+            span.parent_span_id = None;
+        }
+
         span
     }
 
@@ -504,6 +514,10 @@ fn should_keep_attribute(attribute: &str) -> bool {
         return false;
     }
 
+    if attribute == OVERRIDE_PARENT_SPAN_ATTRIBUTE_NAME {
+        return false;
+    }
+
     // remove gen_ai.prompt/completion attributes as they are stored in LLM span's input/output
     let pattern = Regex::new(r"gen_ai\.(prompt|completion)\.\d+\.(content|role)").unwrap();
     return !pattern.is_match(attribute);