From 961fb3f2b554c44a64cf9e75be8ef3d364a05c59 Mon Sep 17 00:00:00 2001
From: Dinmukhamed Mailibay <47117969+dinmukhamedm@users.noreply.github.com>
Date: Mon, 18 Nov 2024 18:56:48 -0800
Subject: [PATCH] treat more spans as llm spans (#222)

---
 app-server/src/traces/spans.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/app-server/src/traces/spans.rs b/app-server/src/traces/spans.rs
index 004ffff1..d766d1aa 100644
--- a/app-server/src/traces/spans.rs
+++ b/app-server/src/traces/spans.rs
@@ -141,7 +141,12 @@ impl SpanAttributes {
             serde_json::from_value::<SpanType>(span_type.clone()).unwrap_or_default()
         } else {
             // quick hack until we figure how to set span type on auto-instrumentation
-            if self.attributes.contains_key(GEN_AI_SYSTEM) {
+            if self.attributes.contains_key(GEN_AI_SYSTEM)
+                || self
+                    .attributes
+                    .iter()
+                    .any(|(k, _)| k.starts_with("gen_ai.") || k.starts_with("llm."))
+            {
                 SpanType::LLM
             } else {
                 SpanType::DEFAULT