treat more spans as llm spans (#222)

lmnr-ai · Nov 19, 2024 · 961fb3f · 961fb3f
1 parent 13c7b5b
commit 961fb3f
Showing 1 changed file with 6 additions and 1 deletion.
diff --git 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